network.py

import logging
from pathlib import Path

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

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

logger = logging.getLogger(__name__)


@attr.s
class FormattedData:
    unique_x: np.array = attr.ib()
    probs_x: np.array = attr.ib()
    y: np.array = attr.ib()


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()):
        self.model = nn.Sequential(
            nn.Linear(input_features, 10),
            nn.ReLU(),
            nn.Linear(10, 12),
            nn.ReLU(),
            nn.Linear(12, 1),
            nn.Sigmoid()
        )
        self._set_weights()
        self.conf = nn_conf
        self.id = OrderCounter.next()

        # Variables used for loss function
        self.attacker_actions: FormattedData = None
        self.benign_data: FormattedData = None

        # Variables from last training epoch measuring quality
        self.final_loss = None
        self.final_fp_cost = None

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

    def set_data(self, benign_data: FormattedData, attack: FormattedData):
        self.attacker_actions = attack
        self.benign_data = benign_data

    def loss_function(self, x, limits, real_y, probs):
        zero_sum_part = torch.sum(real_y*(1-limits)*torch.prod(x, dim=1)*probs)
        fp_cost = self._fp_cost_tensor(limits, real_y, probs)
        sum_loss = torch.add(zero_sum_part, fp_cost)
        return sum_loss

    def _fp_cost_tensor(self, limits, real_y, probs):
        return torch.sum((1-real_y) * probs * torch.pow(limits, 4))

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

        x = np.concatenate((defender.unique_x, attacker.unique_x), axis=0)
        y = np.concatenate((defender.y, attacker.y), axis=0)
        probs = np.concatenate((defender.probs_x, attacker.probs_x), axis=0)

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

        self.x_train = torch.from_numpy(x).float()
        self.y_train = torch.from_numpy(y).float()
        self.probs_train = torch.from_numpy(probs).float()

    def _set_weights(self):
        def init_weights(m):
            if type(m) == nn.Linear:
                torch.nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(.0)
        self.model.apply(init_weights)

    def train(self):
        self._prepare_data()
        self._train()

    def _train(self):
        learning_rate = self.conf.learning_rate
        optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)

        for e in range(self.conf.epochs):
            # Forward pass: compute predicted y by passing x to the model.
            train_limits = self._limit_predict(self.x_train, with_grad=True)

            # Compute loss.
            loss = self.loss_function(self.x_train, train_limits, self.y_train,
                                      self.probs_train)
            # loss = self.loss_fn(train_limits, self.y_train)

            # Compute validation loss and report some info
            if e % 5 == 0:
                logging.debug(f'Epoch: {e}/{self.conf.epochs},\t'
                              f'TrainLoss: {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()

        self.final_fp_cost = self._fp_cost_tensor(train_limits, self.y_train,
                                                  self.probs_train).item()
        self.final_loss = loss

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

    def _limit_predict(self, x: torch.Tensor, with_grad=False):
        if with_grad:
                raw_prediction = self._raw_predict(x)
        else:
            with torch.no_grad():
                raw_prediction = self._raw_predict(x)

        # The same as lambda p: 0 if p < 0.5 else (p - 0.5) * 2
        # TODO try to use e.g. sigmoid
        clamped = raw_prediction.clamp(min=0.5, max=1)
        limit = torch.mul(torch.add(clamped, -0.5), 2)
        return limit

    def predict_single_limit(self, input, return_tensor=False):
        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()

        if return_tensor:
            return self._limit_predict(input)[0]
        else:
            return self._limit_predict(input)[0].item()


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, _ = np_arrays_from_scored_csv(
        Path('all_benign_scored.csv'), 0, 500)
    malicious_x, _ = np_arrays_from_scored_csv(
        Path('scored_malicious.csv'), 1, 1)

    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))
    benign_data = FormattedData(benign_unique_x, probs_benign, benign_y)

    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))
    malicious_data = FormattedData(malicious_unique_x, probs_malicious, malicious_y)

    nn = NeuralNetwork()
    nn.set_data(benign_data, malicious_data)
    nn.train()