| 4 | |
| 5 | |
| 6 | class TABL_layer(pl.LightningModule): |
| 7 | def __init__(self, d2, d1, t1, t2): |
| 8 | super().__init__() |
| 9 | self.t1 = t1 |
| 10 | |
| 11 | weight = torch.Tensor(d2, d1) |
| 12 | self.W1 = nn.Parameter(weight) |
| 13 | nn.init.kaiming_uniform_(self.W1, nonlinearity='relu') |
| 14 | |
| 15 | weight2 = torch.Tensor(t1, t1) |
| 16 | self.W = nn.Parameter(weight2) |
| 17 | nn.init.constant_(self.W, 1 / t1) |
| 18 | |
| 19 | weight3 = torch.Tensor(t1, t2) |
| 20 | self.W2 = nn.Parameter(weight3) |
| 21 | nn.init.kaiming_uniform_(self.W2, nonlinearity='relu') |
| 22 | |
| 23 | bias1 = torch.Tensor(d2, t2) |
| 24 | self.B = nn.Parameter(bias1) |
| 25 | nn.init.constant_(self.B, 0) |
| 26 | |
| 27 | l = torch.Tensor(1, ) |
| 28 | self.l = nn.Parameter(l) |
| 29 | nn.init.constant_(self.l, 0.5) |
| 30 | |
| 31 | self.activation = nn.ReLU() |
| 32 | |
| 33 | def forward(self, X): |
| 34 | |
| 35 | # maintaining the weight parameter between 0 and 1. |
| 36 | if (self.l[0] < 0): |
| 37 | l = torch.Tensor(1, ) |
| 38 | self.l = nn.Parameter(l) |
| 39 | nn.init.constant_(self.l, 0.0) |
| 40 | |
| 41 | if (self.l[0] > 1): |
| 42 | l = torch.Tensor(1, ) |
| 43 | self.l = nn.Parameter(l) |
| 44 | nn.init.constant_(self.l, 1.0) |
| 45 | |
| 46 | # modelling the dependence along the first mode of X while keeping the temporal order intact (7) |
| 47 | X = self.W1 @ X |
| 48 | |
| 49 | # enforcing constant (1) on the diagonal |
| 50 | W = self.W - self.W * torch.eye(self.t1, dtype=torch.float32, device="cuda") + torch.eye(self.t1, dtype=torch.float32, device="cuda") / self.t1 |
| 51 | |
| 52 | # attention, the aim of the second step is to learn how important the temporal instances are to each other (8) |
| 53 | E = X @ W |
| 54 | |
| 55 | # computing the attention mask (9) |
| 56 | A = torch.softmax(E, dim=-1) |
| 57 | |
| 58 | # applying a soft attention mechanism (10) |
| 59 | # he attention mask A obtained from the third step is used to zero out the effect of unimportant elements |
| 60 | X = self.l[0] * (X) + (1.0 - self.l[0]) * X * A |
| 61 | |
| 62 | # the final step of the proposed layer estimates the temporal mapping W2, after the bias shift (11) |
| 63 | y = X @ self.W2 + self.B |