network.py 7.64 KB
Newer Older
1
import itertools
Martin Řepa's avatar
backup  
Martin Řepa committed
2
import logging
3
import time
4 5
from pathlib import Path

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

14
from config import NeuralNetworkConfig, RootConfig
15
from data.loader import np_arrays_from_scored_csv
16

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

19

20
@attr.s
21
class FormattedData:
22 23 24 25 26
    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
27 28 29 30 31 32 33 34
class OrderCounter:
    order = 0

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

35

36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
class SoftClip(nn.Module):
    """ SoftClipping activation function

    https://arxiv.org/pdf/1810.11509.pdf
    """
    def __init__(self, p=50.0):
        super().__init__()
        self.p = p

    def forward(self, x):
        first_pow = torch.pow(np.e, torch.mul(x, self.p))
        second_pow = torch.pow(np.e, torch.mul(torch.add(x, -1.0), self.p))

        first_div = torch.add(first_pow, 1.0)
        second_div = torch.add(second_pow, 1.0)

        division = torch.div(first_div, second_div)

        second_part_log = torch.log(division)
        first_part = 1.0 / self.p

        if len(second_part_log[torch.isnan(second_part_log)]):
            print("prdel")  # todo remove me

        return torch.mul(second_part_log, first_part)


63
class NeuralNetwork:
64
    def __init__(self, input_features=2,
65
                 nn_conf: NeuralNetworkConfig = NeuralNetworkConfig()):
Martin Řepa's avatar
backup  
Martin Řepa committed
66
        self.model = nn.Sequential(
67 68 69 70 71 72 73
            nn.Linear(input_features, 40),
            nn.LeakyReLU(),
            nn.Linear(40, 30),
            nn.LeakyReLU(),
            nn.Linear(30, 1),
            # nn.Tanh(),
            # SoftClip(50)
Martin Řepa's avatar
backup  
Martin Řepa committed
74
            nn.Sigmoid()
75
        )
76 77
        self._set_weights()
        self.conf = nn_conf
78
        self.id = OrderCounter.next()
79

80
        # Variables used for loss function
81 82 83
        self.attacker_actions: FormattedData = None
        self.benign_data: FormattedData = None

84
        # Variables from last training epoch measuring quality
85
        self.final_loss = None
86
        self.final_fp_cost = None
87

88
    def __str__(self):
89
        return f'Neural network id:{self.id}, final loss: {self.final_loss}'
90

91
    def set_data(self, benign_data: FormattedData, attack: FormattedData):
92 93 94 95 96 97
        self.attacker_actions = attack
        self.benign_data = benign_data

    def _prepare_data(self):
        defender = self.benign_data
        attacker = self.attacker_actions
98
        x = np.concatenate((defender.unique_x, attacker.unique_x), axis=0)
99
        y = np.concatenate((defender.y, attacker.y), axis=0)
100
        probs = np.concatenate((defender.probs_x, attacker.probs_x), axis=0)
101 102 103 104

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

105 106 107
        self.x_train = torch.tensor(x).float()
        self.y_train = torch.tensor(y).float()
        self.probs_train = torch.tensor(probs).float()
108

109 110 111
    def _set_weights(self):
        def init_weights(m):
            if type(m) == nn.Linear:
112
                torch.nn.init.xavier_uniform_(m.weight)
113 114 115
                m.bias.data.fill_(.0)
        self.model.apply(init_weights)

116 117 118 119 120
    def train(self):
        self._prepare_data()
        self._train()

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

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

Martin Řepa's avatar
Martin Řepa committed
128 129 130
            # 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
131

Martin Řepa's avatar
Martin Řepa committed
132
            # Log loss function value each 5 epochs
133
            if e % 50 == 0:
134
                logging.debug(f'Epoch: {e}/{self.conf.epochs},\t'
135
                              f'TrainLoss: {loss},\t')
Martin Řepa's avatar
backup  
Martin Řepa committed
136 137 138 139

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

Martin Řepa's avatar
backup  
Martin Řepa committed
141 142 143
            # Backward pass: compute gradient of the loss with respect to model
            # parameters
            loss.backward()
144

Martin Řepa's avatar
backup  
Martin Řepa committed
145 146 147
            # Calling the step function on an Optimizer makes an update to its
            # parameters
            optimizer.step()
148

Martin Řepa's avatar
Martin Řepa committed
149
        with torch.no_grad():
150 151
            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
152 153 154
            # measuring quality of final network
            self.final_loss = loss.item()
            self.final_fp_cost = fp_part.item()
155

156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179
    def plot_tmp(self):
        # todo remove me
        if self.id == 1:
            return
        try:
            actions = self.actions
        except AttributeError:
            self.plotted = []
            self.plotting = plt.subplots()
            one_axis = np.linspace(0, 1, 101)  # [0.00, 0.01, 0.02, ..., 0.99, 1.00]
            generator = itertools.product(*itertools.repeat(one_axis, 2))
            actions = torch.tensor(np.array(list(generator))).float()
            self.actions = actions
        finally:
            # Remove all lines from previous iteration plotting
            for item in self.plotted:
                item.remove()
            self.plotted = []
            res = self.latency_predict(actions).numpy().reshape((101, 101), order='F')
            self.plotted.append(self.plotting[1].imshow(res, cmap='Reds', vmin=0, vmax=1, origin='lower'))
            self.plotting[0].canvas.draw()
            plt.pause(0.000001)
            time.sleep(1)

180
    def _raw_predict(self, tensor: torch.Tensor):
181 182
        pred = self.model(tensor)
        return pred.flatten().float()
183

184
    def latency_predict(self, x: torch.Tensor, with_grad=False):
185 186 187 188 189
        if with_grad:
                raw_prediction = self._raw_predict(x)
        else:
            with torch.no_grad():
                raw_prediction = self._raw_predict(x)
190

191
        return raw_prediction
Martin Řepa's avatar
backup  
Martin Řepa committed
192

Martin Řepa's avatar
Martin Řepa committed
193
    def predict_single_latency(self, input, return_tensor=False):
194 195 196 197 198
        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()

199
        if return_tensor:
200
            return self.latency_predict(input)[0]
201
        else:
202
            return self.latency_predict(input)[0].item()
203

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

Martin Řepa's avatar
Martin Řepa committed
205
def setup_loger(debug: bool):
Martin Řepa's avatar
backup  
Martin Řepa committed
206 207
    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
208
    level = logging.DEBUG if debug else logging.INFO
Martin Řepa's avatar
backup  
Martin Řepa committed
209 210 211 212
    logging.basicConfig(level=level, format=log_format)


if __name__ == '__main__':
Martin Řepa's avatar
Martin Řepa committed
213
    setup_loger(True)
214
    benign_x, _ = np_arrays_from_scored_csv(
215
        Path('all_benign_scored.csv'), 0, 500)
216
    malicious_x, _ = np_arrays_from_scored_csv(
217
        Path('scored_malicious.csv'), 1, 400)
218 219 220 221

    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))
222
    benign_data = FormattedData(benign_unique_x, probs_benign, benign_y)
223 224 225 226

    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))
227
    malicious_data = FormattedData(malicious_unique_x, probs_malicious, malicious_y)
Martin Řepa's avatar
backup  
Martin Řepa committed
228

Martin Řepa's avatar
Martin Řepa committed
229
    conf = RootConfig()
230
    nn = NeuralNetwork(2, conf.model_conf.defender_conf.nn_conf)
231 232
    nn.set_data(benign_data, malicious_data)
    nn.train()