network.py 5.86 KB
Newer Older
Martin Řepa's avatar
backup  
Martin Řepa committed
1
import logging
2
from pathlib import Path
Martin Řepa's avatar
Martin Řepa committed
3
from typing import Callable
4

5
import attr
6
import numpy as np
Martin Řepa's avatar
backup  
Martin Řepa committed
7
import torch
8
from sklearn.utils import shuffle
Martin Řepa's avatar
backup  
Martin Řepa committed
9 10
from torch import nn
from torch import optim
11

12
from config import NeuralNetworkConfig, RootConfig
13 14
from src.data.loader import np_arrays_from_scored_csv

Martin Řepa's avatar
backup  
Martin Řepa committed
15 16
logger = logging.getLogger(__name__)

17

18
@attr.s
19
class FormattedData:
20 21 22 23 24
    unique_x: np.array = attr.ib()
    probs_x: np.array = attr.ib()
    y: np.array = attr.ib()


Martin Řepa's avatar
backup  
Martin Řepa committed
25 26 27 28 29 30 31 32
class OrderCounter:
    order = 0

    @staticmethod
    def next():
        OrderCounter.order += 1
        return OrderCounter.order

33 34

class NeuralNetwork:
35
    def __init__(self, input_features=2,
36
                 nn_conf: NeuralNetworkConfig = NeuralNetworkConfig()):
Martin Řepa's avatar
backup  
Martin Řepa committed
37 38 39 40 41 42 43
        self.model = nn.Sequential(
            nn.Linear(input_features, 10),
            nn.ReLU(),
            nn.Linear(10, 12),
            nn.ReLU(),
            nn.Linear(12, 1),
            nn.Sigmoid()
44
        )
45 46
        self._set_weights()
        self.conf = nn_conf
47
        self.id = OrderCounter.next()
48

49
        # Variables used for loss function
50 51 52
        self.attacker_actions: FormattedData = None
        self.benign_data: FormattedData = None

53
        # Variables from last training epoch measuring quality
54
        self.final_loss = None
55
        self.final_fp_cost = None
56

57
    def __str__(self):
58
        return f'Neural network id:{self.id}, final loss: {self.final_loss}'
59

60
    def set_data(self, benign_data: FormattedData, attack: FormattedData):
61 62 63 64 65 66 67
        self.attacker_actions = attack
        self.benign_data = benign_data

    def _prepare_data(self):
        defender = self.benign_data
        attacker = self.attacker_actions

68
        x = np.concatenate((defender.unique_x, attacker.unique_x), axis=0)
69
        y = np.concatenate((defender.y, attacker.y), axis=0)
70
        probs = np.concatenate((defender.probs_x, attacker.probs_x), axis=0)
71 72 73 74

        # Shuffle before splitting
        x, y, probs = shuffle(x, y, probs, random_state=1)

75 76 77
        self.x_train = torch.from_numpy(x).float()
        self.y_train = torch.from_numpy(y).float()
        self.probs_train = torch.from_numpy(probs).float()
78

79 80 81
    def _set_weights(self):
        def init_weights(m):
            if type(m) == nn.Linear:
82
                torch.nn.init.xavier_uniform_(m.weight)
83 84 85
                m.bias.data.fill_(.0)
        self.model.apply(init_weights)

86 87 88 89 90
    def train(self):
        self._prepare_data()
        self._train()

    def _train(self):
91
        learning_rate = self.conf.learning_rate
Martin Řepa's avatar
backup  
Martin Řepa committed
92
        optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)
93

94
        for e in range(self.conf.epochs):
Martin Řepa's avatar
Martin Řepa committed
95
            # Forward pass: compute predicted y by passing x to the model
96
            train_ltncies = self.latency_predict(self.x_train, with_grad=True)
97

Martin Řepa's avatar
Martin Řepa committed
98 99 100
            # Compute loss
            loss, _ = self.conf.loss_function(self.x_train, train_ltncies,
                                              self.y_train, self.probs_train)
Martin Řepa's avatar
backup  
Martin Řepa committed
101

Martin Řepa's avatar
Martin Řepa committed
102
            # Log loss function value each 5 epochs
Martin Řepa's avatar
backup  
Martin Řepa committed
103
            if e % 5 == 0:
104
                logging.debug(f'Epoch: {e}/{self.conf.epochs},\t'
105
                              f'TrainLoss: {loss},\t')
Martin Řepa's avatar
backup  
Martin Řepa committed
106 107 108 109

            # Before the backward pass, use the optimizer object to zero all of
            # the gradients for the variables it will update
            optimizer.zero_grad()
110

Martin Řepa's avatar
backup  
Martin Řepa committed
111 112 113
            # Backward pass: compute gradient of the loss with respect to model
            # parameters
            loss.backward()
114

Martin Řepa's avatar
backup  
Martin Řepa committed
115 116 117
            # Calling the step function on an Optimizer makes an update to its
            # parameters
            optimizer.step()
118

Martin Řepa's avatar
Martin Řepa committed
119
        with torch.no_grad():
120 121
            loss, fp_part = self.conf.loss_function(self.x_train, train_ltncies,
                                                 self.y_train, self.probs_train)
Martin Řepa's avatar
Martin Řepa committed
122 123 124
            # measuring quality of final network
            self.final_loss = loss.item()
            self.final_fp_cost = fp_part.item()
125

126
    def _raw_predict(self, tensor: torch.Tensor):
127 128
        pred = self.model(tensor)
        return pred.flatten().float()
129

130
    def latency_predict(self, x: torch.Tensor, with_grad=False):
131 132 133 134 135
        if with_grad:
                raw_prediction = self._raw_predict(x)
        else:
            with torch.no_grad():
                raw_prediction = self._raw_predict(x)
136

137
        return raw_prediction
Martin Řepa's avatar
backup  
Martin Řepa committed
138

Martin Řepa's avatar
Martin Řepa committed
139
    def predict_single_latency(self, input, return_tensor=False):
140 141 142 143 144
        in_type = type(input)
        if in_type == list or in_type == tuple or \
                in_type == np.array or in_type == np.ndarray:
            input = torch.tensor(input).float()

145
        if return_tensor:
146
            return self.latency_predict(input)[0]
147
        else:
148
            return self.latency_predict(input)[0].item()
149

Martin Řepa's avatar
backup  
Martin Řepa committed
150

Martin Řepa's avatar
Martin Řepa committed
151
def setup_loger(debug: bool):
Martin Řepa's avatar
backup  
Martin Řepa committed
152 153
    log_format = ('%(asctime)-15s\t%(name)s:%(levelname)s\t'
                  '%(module)s:%(funcName)s:%(lineno)s\t%(message)s')
Martin Řepa's avatar
Martin Řepa committed
154
    level = logging.DEBUG if debug else logging.INFO
Martin Řepa's avatar
backup  
Martin Řepa committed
155 156 157 158
    logging.basicConfig(level=level, format=log_format)


if __name__ == '__main__':
Martin Řepa's avatar
Martin Řepa committed
159
    setup_loger(True)
160
    benign_x, _ = np_arrays_from_scored_csv(
161
        Path('all_benign_scored.csv'), 0, 500)
162
    malicious_x, _ = np_arrays_from_scored_csv(
163
        Path('scored_malicious.csv'), 1, 400)
164 165 166 167

    benign_unique_x, counts = np.unique(benign_x, axis=0, return_counts=True)
    probs_benign = np.array([count / len(benign_x) for count in counts])
    benign_y = np.zeros(len(benign_unique_x))
168
    benign_data = FormattedData(benign_unique_x, probs_benign, benign_y)
169 170 171 172

    malicious_unique_x, counts = np.unique(malicious_x, axis=0, return_counts=True)
    probs_malicious = np.array([count / len(malicious_unique_x) for count in counts])
    malicious_y = np.ones(len(malicious_unique_x))
173
    malicious_data = FormattedData(malicious_unique_x, probs_malicious, malicious_y)
Martin Řepa's avatar
backup  
Martin Řepa committed
174

Martin Řepa's avatar
Martin Řepa committed
175
    conf = RootConfig()
176
    nn = NeuralNetwork(2, conf.model_conf.defender_conf.nn_conf)
177 178
    nn.set_data(benign_data, malicious_data)
    nn.train()
Martin Řepa's avatar
backup  
Martin Řepa committed
179