| 9 | |
| 10 | class MLP(torch.nn.Module): |
| 11 | def __init__(self, input_size: int, xcol: str = "emb", ycol: str = "avg_rating"): |
| 12 | super().__init__() |
| 13 | self.input_size = input_size |
| 14 | self.xcol = xcol |
| 15 | self.ycol = ycol |
| 16 | self.layers = torch.nn.Sequential( |
| 17 | torch.nn.Linear(self.input_size, 1024), |
| 18 | #torch.nn.ReLU(), |
| 19 | torch.nn.Dropout(0.2), |
| 20 | torch.nn.Linear(1024, 128), |
| 21 | #torch.nn.ReLU(), |
| 22 | torch.nn.Dropout(0.2), |
| 23 | torch.nn.Linear(128, 64), |
| 24 | #torch.nn.ReLU(), |
| 25 | torch.nn.Dropout(0.1), |
| 26 | torch.nn.Linear(64, 16), |
| 27 | #torch.nn.ReLU(), |
| 28 | torch.nn.Linear(16, 1), |
| 29 | ) |
| 30 | |
| 31 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 32 | return self.layers(x) |