<filename>src/neural_network/convolutional_neural_network.py
import numpy
import torch.nn as nn

from neural_network.abstract_neural_network import AbstractNeuralNetwork


class ConvolutionalNeuralNetwork(AbstractNeuralNetwork):

    def __init__(self, device):
        super(ConvolutionalNeuralNetwork, self).__init__(device)
        # building nn
        self.convolutional_layer = nn.Sequential(
            # Conv Layer block 1
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # Conv Layer block 2
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Dropout2d(p=0.05),
            # Conv Layer block 3
            nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.fully_connected_layer = nn.Sequential(
            nn.Dropout(p=0.1),
            nn.Linear(4096, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 512),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.1),
            nn.Linear(512, 10)
        )
        self.device = device
        self.to(device)

    def forward(self, x):
        x = x.to(self.device)
        x = self.convolutional_layer(x)
        x = x.view(x.size(0), -1)
        x = self.fully_connected_layer(x)
        return x

    def _get_gradient(self):
        self.to("cpu")
        gradient = numpy.concatenate([self.convolutional_layer[0].weight.grad[2].numpy().flatten(),
                                     self.convolutional_layer[0].weight.grad[4].numpy().flatten(),
                                     self.convolutional_layer[0].weight.grad[8].numpy().flatten(),
                                     self.convolutional_layer[0].weight.grad[10].numpy().flatten(),
                                     self.convolutional_layer[0].weight.grad[15].numpy().flatten(),
                                     self.convolutional_layer[0].weight.grad[17].numpy().flatten()])
        self.to(self.device)
        return gradient

