network.py

import logging
from pathlib import Path
from typing import List, Tuple

import numpy as np
import torch
from sklearn.model_selection import train_test_split
from sklearn.utils import shuffle
from torch import nn
from torch import optim

from config import NeuralNetworkConfig, TrainingNnConfig, RootConfig
from src.data.loader import np_arrays_from_scored_csv

logger = logging.getLogger(__name__)


class OrderCounter:
    order = 0

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


class NeuralNetwork:
    def __init__(self, input_features=2,
                 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()
        )
        self.loss_fn = nn.BCELoss()
        self.attacker_actions = None
        self.epochs = nn_conf.epochs
        self.validation_split = nn_train_conf.validation_split

        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]):
        x, y = benign_data

        # Add attacker's malicious actions to dataset
        attacker_features_x = np.array(attacker_features_x)
        if len(attacker_features_x[0]):
            attacker_features_y = [[1] for _ in attacker_features_x]
            x = np.concatenate((x, attacker_features_x), axis=0)
            y = np.concatenate((y, attacker_features_y), axis=0)

        # Shuffle benign and malicious data
        x, y = shuffle(x, y, random_state=1)

        # 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))

        # Return final data (x_train, x_validate, y_train, y_validate)
        return data

    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)

    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

        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)

            # 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()

            # Backward pass: compute gradient of the loss with respect to model
            # parameters
            loss.backward()

            # Calling the step function on an Optimizer makes an update to its
            # parameters
            optimizer.step()

    def raw_predict(self, x):
        with torch.no_grad():
            tensor = torch.tensor(x).float()
            res = self.model(tensor)
        return res.numpy()

    def limit_predict(self, x):
        raw_prediction = self.raw_predict(x)

        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())