|
|
|
|
|
|
|
|
|
|
|
from torch import nn |
|
|
|
|
|
class MLP(nn.Module): |
|
|
""" |
|
|
Multi layer perceptron. |
|
|
""" |
|
|
def __init__(self, size_in, size_out, size_hidden=None, dropout=0.0): |
|
|
super().__init__() |
|
|
if size_hidden is None: |
|
|
size_hidden = [] |
|
|
sizes = [size_in] + size_hidden + [size_out] |
|
|
|
|
|
net = [] |
|
|
for i in range(len(sizes) - 2): |
|
|
net.append(nn.Linear(sizes[i], sizes[i+1])) |
|
|
net.append(nn.ReLU()) |
|
|
net.append(nn.Dropout(dropout)) |
|
|
|
|
|
net.append(nn.Linear(sizes[-2], sizes[-1])) |
|
|
net = nn.Sequential(*net) |
|
|
self.net = net |
|
|
|
|
|
def forward(self, x): |
|
|
""" |
|
|
Forward method. |
|
|
""" |
|
|
x = self.net(x) |
|
|
return x |
|
|
|
