| 52 | |
| 53 | |
| 54 | class Transformer(pl.LightningModule): |
| 55 | def __init__( |
| 56 | self, |
| 57 | lighten, |
| 58 | dropout: float = 0.1, |
| 59 | activation: str = "relu", |
| 60 | norm_first: bool = False, |
| 61 | ): |
| 62 | super().__init__() |
| 63 | self.name = "transformer" |
| 64 | if lighten: |
| 65 | self.name += "-lighten" |
| 66 | |
| 67 | d_model = 64 if not lighten else 32 |
| 68 | dim_feedforward = 256 if not lighten else 128 |
| 69 | nhead = 8 if not lighten else 4 |
| 70 | num_layers = 2 if not lighten else 1 |
| 71 | |
| 72 | self.embed = nn.Linear(40, d_model, bias=False) |
| 73 | |
| 74 | self.embed_positions = SinusoidalPositionalEmbedding(100, d_model) |
| 75 | |
| 76 | layer_norm_eps: float = 1e-5 |
| 77 | encoder_layer = nn.TransformerEncoderLayer( |
| 78 | d_model=d_model, |
| 79 | nhead=nhead, |
| 80 | dim_feedforward=dim_feedforward, |
| 81 | dropout=dropout, |
| 82 | activation=activation, |
| 83 | layer_norm_eps=layer_norm_eps, |
| 84 | norm_first=norm_first, |
| 85 | batch_first=True, |
| 86 | ) |
| 87 | encoder_norm = nn.LayerNorm(d_model, eps=layer_norm_eps) |
| 88 | self.transformer_encoder = nn.TransformerEncoder( |
| 89 | encoder_layer, num_layers=num_layers, norm=encoder_norm |
| 90 | ) |
| 91 | self.cat_head = nn.Linear(d_model, 3) |
| 92 | |
| 93 | def forward(self, x): |
| 94 | x = self.embed(x.squeeze(1)) |
| 95 | |
| 96 | embed_pos = self.embed_positions(x.shape) |
| 97 | |
| 98 | # transformer encoder |
| 99 | x = self.transformer_encoder(x + embed_pos) |
| 100 | |
| 101 | # mean pool for classification |
| 102 | x = torch.mean(x, dim=1) |
| 103 | |
| 104 | logits = self.cat_head(x) |
| 105 | return logits |