A LayerNorm variant, popularized by Transformers, that performs point-wise mean and variance normalization over the channel dimension for inputs that have shape (batch_size, channels, height, width). https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc9
| 11 | from core.utils import NestedTensor |
| 12 | |
| 13 | class Norm2d(nn.Module): |
| 14 | """ |
| 15 | A LayerNorm variant, popularized by Transformers, that performs point-wise mean and |
| 16 | variance normalization over the channel dimension for inputs that have shape |
| 17 | (batch_size, channels, height, width). |
| 18 | https://github.com/facebookresearch/ConvNeXt/blob/d1fa8f6fef0a165b27399986cc2bdacc92777e40/models/convnext.py#L119 # noqa B950 |
| 19 | """ |
| 20 | |
| 21 | def __init__(self, embed_dim, eps=1e-6): |
| 22 | super().__init__() |
| 23 | self.weight = nn.Parameter(torch.ones(embed_dim)) |
| 24 | self.bias = nn.Parameter(torch.zeros(embed_dim)) |
| 25 | self.eps = eps |
| 26 | self.normalized_shape = (embed_dim,) |
| 27 | |
| 28 | # >>> workaround for compatability |
| 29 | self.ln = nn.LayerNorm(embed_dim, eps=1e-6) |
| 30 | self.ln.weight = self.weight |
| 31 | self.ln.bias = self.bias |
| 32 | |
| 33 | def forward(self, x): |
| 34 | u = x.mean(1, keepdim=True) |
| 35 | s = (x - u).pow(2).mean(1, keepdim=True) |
| 36 | x = (x - u) / torch.sqrt(s + self.eps) |
| 37 | x = self.weight[:, None, None] * x + self.bias[:, None, None] |
| 38 | return x |
| 39 | |
| 40 | |
| 41 | class Conv2d(torch.nn.Conv2d): |