Spaces:
Running on Zero
Running on Zero
File size: 1,146 Bytes
7405e85 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | import torch.nn as nn
def zero_module(module):
"""Zero out a module's parameters and return it."""
for param in module.parameters():
param.detach().zero_()
return module
def normalization(channels, channels_per_group=None):
if channels_per_group is not None:
return Float32GroupNorm(channels // channels_per_group, channels)
num_groups = 4 if channels % 8 != 0 else 8
return Float32GroupNorm(num_groups, channels)
class Float32GroupNorm(nn.GroupNorm):
def forward(self, x):
return super().forward(x.float()).to(dtype=x.dtype)
def conv_nd(dims, *args, **kwargs):
if dims == 1:
return nn.Conv1d(*args, **kwargs)
if dims == 2:
return nn.Conv2d(*args, **kwargs)
if dims == 3:
return nn.Conv3d(*args, **kwargs)
raise ValueError(f"unsupported dimensions: {dims}")
def avg_pool_nd(dims, *args, **kwargs):
if dims == 1:
return nn.AvgPool1d(*args, **kwargs)
if dims == 2:
return nn.AvgPool2d(*args, **kwargs)
if dims == 3:
return nn.AvgPool3d(*args, **kwargs)
raise ValueError(f"unsupported dimensions: {dims}")
|