| import torch |
| import torch.nn as nn |
|
|
| class VanillaMLP(nn.Module): |
| def __init__(self, input_dim, output_dim, out_activation, n_hidden_layers=4, n_neurons=64, activation="ReLU"): |
| super().__init__() |
| self.n_neurons = n_neurons |
| self.n_hidden_layers = n_hidden_layers |
| self.activation = activation |
| self.out_activation = out_activation |
| layers = [ |
| self.make_linear(input_dim, self.n_neurons, is_first=True, is_last=False), |
| self.make_activation(), |
| ] |
| for i in range(self.n_hidden_layers - 1): |
| layers += [ |
| self.make_linear( |
| self.n_neurons, self.n_neurons, is_first=False, is_last=False |
| ), |
| self.make_activation(), |
| ] |
| layers += [ |
| self.make_linear(self.n_neurons, output_dim, is_first=False, is_last=True) |
| ] |
| if self.out_activation == "sigmoid": |
| layers += [nn.Sigmoid()] |
| elif self.out_activation == "tanh": |
| layers += [nn.Tanh()] |
| elif self.out_activation == "hardtanh": |
| layers += [nn.Hardtanh()] |
| elif self.out_activation == "GELU": |
| layers += [nn.GELU()] |
| elif self.out_activation == "RELU": |
| layers += [nn.ReLU()] |
| else: |
| raise NotImplementedError |
| self.layers = nn.Sequential(*layers) |
|
|
| def forward(self, x, split_size=100000): |
| with torch.cuda.amp.autocast(enabled=False): |
| out = self.layers(x) |
| return out |
|
|
| def make_linear(self, dim_in, dim_out, is_first, is_last): |
| layer = nn.Linear(dim_in, dim_out, bias=False) |
| return layer |
|
|
| def make_activation(self): |
| if self.activation == "ReLU": |
| return nn.ReLU(inplace=True) |
| elif self.activation == "GELU": |
| return nn.GELU() |
| else: |
| raise NotImplementedError |
