Spaces:
Running on Zero
Running on Zero
| 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}") | |