| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| class DeepGBM(nn.Module): |
| def __init__(self, num_features, cat_features, num_classes, d_main=128, d_hidden=64, n_blocks=4, dropout=0.2): |
| super(DeepGBM, self).__init__() |
|
|
| self.num_classes = num_classes |
| |
| |
| self.num_linear = nn.Linear(num_features, d_main) |
|
|
| |
| self.cat_embedding = nn.ModuleList([ |
| nn.Embedding(cat_size, d_main) for cat_size in cat_features |
| ]) |
|
|
| |
| self.blocks = nn.ModuleList([ |
| nn.Sequential( |
| nn.Linear(d_main, d_hidden), |
| nn.ReLU(), |
| nn.Dropout(dropout), |
| nn.Linear(d_hidden, d_main), |
| nn.ReLU() |
| ) 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_num = self.num_linear(x_num) |
|
|
| |
| x_cat = [embed(x_cat[:, i]) for i, embed in enumerate(self.cat_embedding)] |
| x_cat = torch.stack(x_cat, dim=1).sum(dim=1) |
| x = x_num + x_cat |
| |
| for block in self.blocks: |
| x = x + block(x) |
| x = self.output_layer(x) |
| |
| return x |
|
|
