File size: 815 Bytes
c176aea |
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 |
from torch import nn
class Net(nn.Module):
"""4 layer model, different activations and neurons count on layer
"""
def __init__(self,input_dim:int=2,hidden_dim:int=200):
"""Init
Args:
input_dim (int, optional): Defaults to 2.
hidden_dim (int, optional): Defaults to 200.
"""
super(Net,self).__init__()
self.input = nn.Linear(input_dim,40)
self.act1 = nn.Tanh()
self.layer = nn.Linear(40,80)
self.act2 = nn.ReLU()
self.layer1 = nn.Linear(80,hidden_dim)
self.act3 = nn.ReLU()
self.layer2 = nn.Linear(hidden_dim,1)
def forward(self, x):
x = self.act2(self.layer(self.act1(self.input(x))))
x = self.act3(self.layer1(x))
x = self.layer2(x)
return x
|