| 2 | from torch import nn |
| 3 | |
| 4 | class MLP(nn.Module): |
| 5 | def __init__(self, input_dim, hidden_dims, output_dim, dropout): |
| 6 | super(MLP, self).__init__() |
| 7 | layers = list() |
| 8 | curr_dim = input_dim |
| 9 | for hidden_dim in hidden_dims: |
| 10 | layers.append(nn.Linear(curr_dim, hidden_dim)) |
| 11 | layers.append(nn.BatchNorm1d(hidden_dim)) |
| 12 | layers.append(nn.ReLU()) |
| 13 | layers.append(nn.Dropout(p=dropout)) |
| 14 | curr_dim = hidden_dim |
| 15 | layers.append(nn.Linear(curr_dim, output_dim)) |
| 16 | self.mlp = nn.Sequential(*layers) |
| 17 | |
| 18 | def forward(self, input): |
| 19 | return self.mlp(input) |
| 20 | |
| 21 | class MaskAvg(nn.Module): |
| 22 | def __init__(self): |