| 4 | |
| 5 | |
| 6 | class BL_layer(pl.LightningModule): |
| 7 | def __init__(self, d2, d1, t1, t2): |
| 8 | super().__init__() |
| 9 | weight1 = torch.Tensor(d2, d1) |
| 10 | self.W1 = nn.Parameter(weight1) |
| 11 | nn.init.kaiming_uniform_(self.W1, nonlinearity='relu') |
| 12 | |
| 13 | weight2 = torch.Tensor(t1, t2) |
| 14 | self.W2 = nn.Parameter(weight2) |
| 15 | nn.init.kaiming_uniform_(self.W2, nonlinearity='relu') |
| 16 | |
| 17 | bias1 = torch.zeros((d2, t2)) |
| 18 | self.B = nn.Parameter(bias1) |
| 19 | nn.init.constant_(self.B, 0) |
| 20 | |
| 21 | self.activation = nn.ReLU() |
| 22 | |
| 23 | def forward(self, x): |
| 24 | |
| 25 | x = self.activation(self.W1 @ x @ self.W2 + self.B) |
| 26 | |
| 27 | return x |