mirrors
/
pytorch-image-models
mirror of https://github.com/huggingface/pytorch-image-models.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update metaformers.py

pull/1647/head
Fredo Guan 2023-01-07 21:19:00 -08:00
parent d17bcb10a5
commit 01f671ed08
1 changed files with 4 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
timm/models/metaformers.py
View File

@ -347,18 +347,16 @@ class LayerNormGeneral(nn.Module):
self.normalized_dim = normalized_dim
self.use_scale = scale
self.use_bias = bias
self.weight = nn.Parameter(torch.ones(affine_shape)) if scale else None
self.bias = nn.Parameter(torch.zeros(affine_shape)) if bias else None
self.weight = nn.Parameter(torch.ones(affine_shape)) if scale else torch.ones(affine_shape)
self.bias = nn.Parameter(torch.zeros(affine_shape)) if bias else torch.zeros(affine_shape)
self.eps = eps
def forward(self, x):
c = x - x.mean(self.normalized_dim, keepdim=True)
s = c.pow(2).mean(self.normalized_dim, keepdim=True)
x = c / torch.sqrt(s + self.eps)
if self.use_scale:
x = x * self.weight
if self.use_bias:
x = x + self.bias
x = x * self.weight
x = x + self.bias
return x