| 4 | |
| 5 | |
| 6 | class DLA(pl.LightningModule): |
| 7 | def __init__(self, lighten, num_snapshots=100, hidden_size=128): |
| 8 | super().__init__() |
| 9 | self.name = "mlp" |
| 10 | num_features = 40 |
| 11 | if lighten: |
| 12 | self.name += "-lighten" |
| 13 | num_features = 20 |
| 14 | |
| 15 | self.W1 = nn.Linear(num_features, num_features, bias=False) |
| 16 | |
| 17 | self.softmax = nn.Softmax(dim=1) |
| 18 | |
| 19 | self.gru = nn.GRU( |
| 20 | input_size=num_features, |
| 21 | hidden_size=hidden_size, |
| 22 | num_layers=2, |
| 23 | batch_first=True, |
| 24 | dropout=0.5 |
| 25 | ) |
| 26 | |
| 27 | self.W2 = nn.Linear(hidden_size, hidden_size, bias=False) |
| 28 | self.W3 = nn.Linear(num_snapshots*hidden_size, 3) |
| 29 | |
| 30 | def forward(self, x): |
| 31 | # x.shape = [batch_size, num_snapshots, num_features] |
| 32 | x = x.squeeze(1) |
| 33 | |
| 34 | X_tilde = self.W1(x) |
| 35 | # alpha.shape = [batch_size, num_snapshots, num_features] |
| 36 | |
| 37 | alpha = self.softmax(X_tilde) |
| 38 | # alpha.shape = [batch_size, num_snapshots, num_features] |
| 39 | |
| 40 | alpha = torch.mean(alpha, dim=2) |
| 41 | # alpha.shape = [batch_size, num_snapshots] |
| 42 | |
| 43 | x_tilde = torch.einsum('ij,ijk->ijk', [alpha, x]) |
| 44 | # x_tilde.shape = [batch_size, num_snapshots, num_features] |
| 45 | |
| 46 | H, _ = self.gru(x_tilde) |
| 47 | # o.shape = [batch_size, num_snapshots, hidden_size] |
| 48 | |
| 49 | H_tilde = self.W2(H) |
| 50 | # o.shape = [batch_size, num_snapshots, hidden_size] |
| 51 | |
| 52 | beta = self.softmax(H_tilde) |
| 53 | # o.shape = [batch_size, num_snapshots, hidden_size] |
| 54 | |
| 55 | beta = torch.mean(beta, dim=2) |
| 56 | # beta.shape = [batch_size, num_snapshots] |
| 57 | |
| 58 | h_tilde = torch.einsum('ij,ijk->ijk', [beta, H]) |
| 59 | # h_tilde.shape = [batch_size, num_snapshots, hidden_size] |
| 60 | |
| 61 | h_tilde = torch.flatten(h_tilde, start_dim=1) |
| 62 | # h_tilde.shape = [batch_size, hidden_size*num_snapshots] |
| 63 | |