model/MLP.py

"""
Definition of the class of models with 2 hidden layers.
"""
import torch
from torch import nn

class MLP_C2H2(nn.Module):
    def __init__(self,
        input_size=1024,
        hidden_size_1=512,
        hidden_size_2=256
    ):
        super().__init__()
        self.input_size = input_size
        self.hidden_size_1 = hidden_size_1
        self.hidden_size_2 = hidden_size_2

        self.model = torch.nn.ModuleList(
            [
                nn.Linear(self.input_size, self.hidden_size_1),
                nn.ReLU(),
            ]
            + [
                nn.Linear(self.hidden_size_1, self.hidden_size_2),
                nn.ReLU(),
            ]
            + [
                nn.Linear(self.hidden_size_2, 1),
                nn.Sigmoid()
            ]
        )
        self.loss_function = nn.BCELoss()

    def forward(self, point):
        for layer in self.model:
            point = layer(point)
        return point

    def calculate_loss(self, point, label):
        return self.loss_function(point, label)