| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| |
| class ResNetLike(nn.Module): |
| def __init__(self, num_classes, input_dim, d_main=128, d_hidden=64, n_blocks=4, dropout_first=0.25, dropout_second=0.0): |
| super(ResNetLike, self).__init__() |
|
|
| self.num_classes = num_classes |
|
|
| |
| self.input_layer = nn.Linear(input_dim, d_main) |
|
|
| |
| self.blocks = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(d_main, d_hidden), |
| nn.ReLU(), |
| nn.Dropout(dropout_first), |
| nn.Linear(d_hidden, d_main), |
| nn.Dropout(dropout_second) |
| ) for _ in range(n_blocks) |
| ]) |
| |
| if num_classes == 2: |
| self.output_layer = nn.Linear(d_main, 1) |
| elif num_classes > 2: |
| self.output_layer = nn.Linear(d_main, num_classes) |
|
|
| def forward(self, x_num, x_cat): |
| x = torch.cat([x_num, x_cat], dim=1) |
| x = F.relu(self.input_layer(x)) |
| |
| for block in self.blocks: |
| x = x + block(x) |
| x = self.output_layer(x) |
| |
| return x |
|
|
