Simple 2-layer MLP for regression
| 18 | |
| 19 | |
| 20 | class SimpleMLP(nn.Module): |
| 21 | """Simple 2-layer MLP for regression""" |
| 22 | def __init__(self, input_dim=20, hidden_dims=[64, 32], dropout=0.0): |
| 23 | super().__init__() |
| 24 | layers = [] |
| 25 | prev_dim = input_dim |
| 26 | for hidden_dim in hidden_dims: |
| 27 | layers.append(nn.Linear(prev_dim, hidden_dim)) |
| 28 | layers.append(nn.ReLU()) |
| 29 | if dropout > 0: |
| 30 | layers.append(nn.Dropout(dropout)) |
| 31 | prev_dim = hidden_dim |
| 32 | layers.append(nn.Linear(prev_dim, 1)) |
| 33 | self.network = nn.Sequential(*layers) |
| 34 | |
| 35 | def forward(self, x): |
| 36 | return self.network(x).squeeze(-1) |
| 37 | |
| 38 | |
| 39 | def generate_data(n_samples=1000, n_features=20, noise=10.0, random_state=42): |