From 8165cacd8280dd61b0ed081c1f7fdf18fa41937e Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Mon, 5 Jul 2021 18:21:34 -0700 Subject: [PATCH] Realized LayerNorm2d won't work in all cases as is, fixed. --- timm/models/layers/norm.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/timm/models/layers/norm.py b/timm/models/layers/norm.py index 433552b4..aace107b 100644 --- a/timm/models/layers/norm.py +++ b/timm/models/layers/norm.py @@ -15,9 +15,10 @@ class GroupNorm(nn.GroupNorm): class LayerNorm2d(nn.LayerNorm): - """ Layernorm for channels of '2d' spatial BCHW tensors """ + """ LayerNorm for channels of '2D' spatial BCHW tensors """ def __init__(self, num_channels): - super().__init__([num_channels, 1, 1]) + super().__init__(num_channels) def forward(self, x: torch.Tensor) -> torch.Tensor: - return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + return F.layer_norm( + x.permute(0, 2, 3, 1), self.normalized_shape, self.weight, self.bias, self.eps).permute(0, 3, 1, 2)