File size: 1,043 Bytes
4b5b766 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """
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)
|