network.py 5.21 KB
Newer Older
Martin Řepa's avatar
backup  
Martin Řepa committed
1
import logging
2 3 4 5
from pathlib import Path
from typing import List, Tuple

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

Martin Řepa's avatar
backup  
Martin Řepa committed
12
from config import NeuralNetworkConfig, TrainingNnConfig, RootConfig
13 14
from src.data.loader import np_arrays_from_scored_csv

Martin Řepa's avatar
backup  
Martin Řepa committed
15 16 17 18 19 20 21 22 23 24 25
logger = logging.getLogger(__name__)


class OrderCounter:
    order = 0

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

26 27

class NeuralNetwork:
28
    def __init__(self, input_features=2,
Martin Řepa's avatar
backup  
Martin Řepa committed
29 30 31 32 33 34 35 36 37
                 nn_conf: NeuralNetworkConfig = NeuralNetworkConfig(),
                 nn_train_conf: TrainingNnConfig = TrainingNnConfig()):
        self.model = nn.Sequential(
            nn.Linear(input_features, 10),
            nn.ReLU(),
            nn.Linear(10, 12),
            nn.ReLU(),
            nn.Linear(12, 1),
            nn.Sigmoid()
38
        )
Martin Řepa's avatar
backup  
Martin Řepa committed
39 40
        self.loss_fn = nn.BCELoss()
        self.attacker_actions = None
41
        self.epochs = nn_conf.epochs
Martin Řepa's avatar
backup  
Martin Řepa committed
42
        self.validation_split = nn_train_conf.validation_split
43

Martin Řepa's avatar
backup  
Martin Řepa committed
44 45 46 47 48 49 50 51 52 53
        self.order = OrderCounter.next()

    def set_attacker_actions(self, attacker_actions: Tuple):
        self.attacker_actions = attacker_actions

    def loss_function(self):
        pass

    def _prepare_data(self, attacker_features_x: List[List[float]],
                      benign_data: Tuple[np.ndarray, np.ndarray]):
54 55
        x, y = benign_data

Martin Řepa's avatar
backup  
Martin Řepa committed
56
        # Add attacker's malicious actions to dataset
57 58
        attacker_features_x = np.array(attacker_features_x)
        if len(attacker_features_x[0]):
Martin Řepa's avatar
backup  
Martin Řepa committed
59
            attacker_features_y = [[1] for _ in attacker_features_x]
60 61 62
            x = np.concatenate((x, attacker_features_x), axis=0)
            y = np.concatenate((y, attacker_features_y), axis=0)

Martin Řepa's avatar
backup  
Martin Řepa committed
63
        # Shuffle benign and malicious data
64 65
        x, y = shuffle(x, y, random_state=1)

Martin Řepa's avatar
backup  
Martin Řepa committed
66 67 68 69 70
        # Split data so we have train dataset and validation dataset
        data = train_test_split(x, y, test_size=self.validation_split)

        # Convert data to float() for pyTorch model compatibility
        data = tuple(map(lambda a: torch.from_numpy(a).float(), data))
71

Martin Řepa's avatar
backup  
Martin Řepa committed
72 73
        # Return final data (x_train, x_validate, y_train, y_validate)
        return data
74

Martin Řepa's avatar
backup  
Martin Řepa committed
75 76 77 78 79 80
    def train(self,
              attacker_features_x: List[List[float]],
              benign_data: Tuple[np.ndarray, np.ndarray]):
        data = self._prepare_data(attacker_features_x, benign_data)
        x_train, x_validate, y_train, y_validate = data
        self._train(x_train, y_train, x_validate, y_validate)
81

Martin Řepa's avatar
backup  
Martin Řepa committed
82 83 84 85
    def _train(self, x, y, x_validate, y_validate):
        learning_rate = 1e-2
        optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)
        x.requires_grad = True
86

Martin Řepa's avatar
backup  
Martin Řepa committed
87 88 89 90
        for e in range(self.epochs):
            logger.debug(f'Running epoch number {e}/{self.epochs}')
            # Forward pass: compute predicted y by passing x to the model.
            y_pred = self.model(x)
91

Martin Řepa's avatar
backup  
Martin Řepa committed
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107
            # Compute and print loss.
            loss = self.loss_fn(y_pred, y)
            logger.debug(f'TestLoss: {loss.item()}, ValidateLoss: todo')  # todo

            # Compute validation loss and report some info
            if e % 5 == 0:
                with torch.no_grad():
                    y_validate_pred = self.model(x_validate)
                    validate_loss = self.loss_fn(y_validate_pred, y_validate)
                logging.debug(f'Epoch: {e}/{self.epochs},\t'
                              f'TrainLoss: {loss.item()},\t'
                              f'ValidateLoss: {validate_loss},\t')

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

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

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

Martin Řepa's avatar
backup  
Martin Řepa committed
117 118 119 120 121
    def raw_predict(self, x):
        with torch.no_grad():
            tensor = torch.tensor(x).float()
            res = self.model(tensor)
        return res.numpy()
122

Martin Řepa's avatar
backup  
Martin Řepa committed
123 124
    def limit_predict(self, x):
        raw_prediction = self.raw_predict(x)
125

Martin Řepa's avatar
backup  
Martin Řepa committed
126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152
        np_limit_func = np.vectorize(lambda p: 0 if p < 0.5 else (p - 0.5) * 2)
        return np_limit_func(raw_prediction)


def setup_loger(conf):
    log_format = ('%(asctime)-15s\t%(name)s:%(levelname)s\t'
                  '%(module)s:%(funcName)s:%(lineno)s\t%(message)s')
    level = logging.DEBUG if conf.base_conf.debug else logging.INFO
    logging.basicConfig(level=level, format=log_format)


if __name__ == '__main__':
    setup_loger(RootConfig())
    benign_x, benign_y = np_arrays_from_scored_csv(
        Path('all_benign_scored.csv'), 0, 1000)
    malicious_x, malicious_y = np_arrays_from_scored_csv(
        Path('scored_malicious.csv'), 1, 0)

    nn = NeuralNetwork()
    nn.train(malicious_x, (benign_x, benign_y))

    # test_loss, test_acc = network.model.evaluate(x_test, y_test)
    # print('Test loss:', test_loss)
    # print('Test accuracy:', test_acc)
    #
    # network.calc_n0_false_positives(benign_x)
    # print(network.get_false_positive_rate())