LayerNorm for channels of '2D' spatial BCHW tensors
| 19 | |
| 20 | |
| 21 | class LayerNorm2d(nn.LayerNorm): |
| 22 | """LayerNorm for channels of '2D' spatial BCHW tensors""" |
| 23 | |
| 24 | def __init__(self, num_channels): |
| 25 | super().__init__(num_channels) |
| 26 | |
| 27 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 28 | return F.layer_norm( |
| 29 | x.permute(0, 2, 3, 1), |
| 30 | self.normalized_shape, |
| 31 | self.weight, |
| 32 | self.bias, |
| 33 | self.eps, |
| 34 | ).permute(0, 3, 1, 2) |
| 35 | |
| 36 | |
| 37 | # model |