| 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 | |