| 84 | |
| 85 | |
| 86 | class MLP(nn.Module): |
| 87 | def __init__(self): |
| 88 | super().__init__() |
| 89 | self.layers = nn.Sequential( |
| 90 | nn.Linear(768, 1024), |
| 91 | nn.Dropout(0.2), |
| 92 | nn.Linear(1024, 128), |
| 93 | nn.Dropout(0.2), |
| 94 | nn.Linear(128, 64), |
| 95 | nn.Dropout(0.1), |
| 96 | nn.Linear(64, 16), |
| 97 | nn.Linear(16, 1), |
| 98 | ) |
| 99 | |
| 100 | @torch.no_grad() |
| 101 | def forward(self, embed): |
| 102 | return self.layers(embed) |
| 103 | |
| 104 | |
| 105 | class AestheticScorer(torch.nn.Module): |