mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix groupnorm specification

Browse Source
This commit is contained in:
James Betker 2021-02-04 22:15:38 -07:00
parent 43da1f9c4b
commit bb79fafb89
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
codes/models/switched_conv/switched_conv_hard_routing.py
View File

@ -185,10 +185,10 @@ class SwitchedConvHardRouting(nn.Module):
self.coupler = Conv2d(coupler_dim_in, breadth, kernel_size=1, stride=self.stride)
elif coupler_mode == 'lambda':
self.coupler = nn.Sequential(nn.Conv2d(coupler_dim_in, coupler_dim_in, 1),
nn.GroupNorm(16, coupler_dim_in),
nn.GroupNorm(2, coupler_dim_in),
nn.ReLU(),
LambdaLayer(dim=coupler_dim_in, dim_out=breadth, r=23, dim_k=16, heads=2, dim_u=1),
nn.GroupNorm(16, breadth),
nn.GroupNorm(1, breadth),
nn.ReLU(),
Conv2d(breadth, breadth, 1, stride=self.stride))
else: