|
|
@ -94,7 +94,7 @@ class LayerNorm2d(nn.LayerNorm):
|
|
|
|
return F.layer_norm(
|
|
|
|
return F.layer_norm(
|
|
|
|
x.permute(0, 2, 3, 1), self.normalized_shape, self.weight, self.bias, self.eps).permute(0, 3, 1, 2)
|
|
|
|
x.permute(0, 2, 3, 1), self.normalized_shape, self.weight, self.bias, self.eps).permute(0, 3, 1, 2)
|
|
|
|
else:
|
|
|
|
else:
|
|
|
|
s, u = torch.var_mean(x, dim=1, keepdim=True)
|
|
|
|
s, u = torch.var_mean(x, dim=1, unbiased=False, keepdim=True)
|
|
|
|
x = (x - u) * torch.rsqrt(s + self.eps)
|
|
|
|
x = (x - u) * torch.rsqrt(s + self.eps)
|
|
|
|
x = x * self.weight[:, None, None] + self.bias[:, None, None]
|
|
|
|
x = x * self.weight[:, None, None] + self.bias[:, None, None]
|
|
|
|
return x
|
|
|
|
return x
|
|
|
|