| 298 | |
| 299 | |
| 300 | class VNLinearBNLeakyReLU(nn.Module): |
| 301 | |
| 302 | def __init__( |
| 303 | self, |
| 304 | in_channels, |
| 305 | out_channels, |
| 306 | dim=5, |
| 307 | share_nonlinearity=False, |
| 308 | negative_slope=0.2, |
| 309 | ): |
| 310 | super().__init__() |
| 311 | |
| 312 | self.linear = VNLinear(in_channels, out_channels) |
| 313 | self.batchnorm = VNBatchNorm(out_channels, dim=dim) |
| 314 | self.leaky_relu = VNLeakyReLU( |
| 315 | out_channels, |
| 316 | # dim=dim, |
| 317 | share_nonlinearity=share_nonlinearity, |
| 318 | negative_slope=negative_slope, |
| 319 | ) |
| 320 | |
| 321 | def forward(self, x): |
| 322 | # Linear |
| 323 | p = self.linear(x) |
| 324 | # BatchNorm |
| 325 | p = self.batchnorm(p) |
| 326 | # LeakyReLU |
| 327 | p = self.leaky_relu(p) |
| 328 | return p |
| 329 | |
| 330 | class NonEquivariantLinearLeakyReLU(nn.Module): |
| 331 | def __init__(self, in_channels, out_channels, dim=2, use_batchnorm=True, negative_slope=0.2): |