| |
| |
| |
| |
|
|
| import torch |
| import torch.nn as nn |
|
|
|
|
| def normalization(channels: int, groups: int = 32): |
| r"""Make a standard normalization layer, i.e. GroupNorm. |
| |
| Args: |
| channels: number of input channels. |
| groups: number of groups for group normalization. |
| |
| Returns: |
| a ``nn.Module`` for normalization. |
| """ |
| assert groups > 0, f"invalid number of groups: {groups}" |
| return nn.GroupNorm(groups, channels) |
|
|
|
|
| def Linear(*args, **kwargs): |
| r"""Wrapper of ``nn.Linear`` with kaiming_normal_ initialization.""" |
| layer = nn.Linear(*args, **kwargs) |
| nn.init.kaiming_normal_(layer.weight) |
| return layer |
|
|
|
|
| def Conv1d(*args, **kwargs): |
| r"""Wrapper of ``nn.Conv1d`` with kaiming_normal_ initialization.""" |
| layer = nn.Conv1d(*args, **kwargs) |
| nn.init.kaiming_normal_(layer.weight) |
| return layer |
|
|
|
|
| def Conv2d(*args, **kwargs): |
| r"""Wrapper of ``nn.Conv2d`` with kaiming_normal_ initialization.""" |
| layer = nn.Conv2d(*args, **kwargs) |
| nn.init.kaiming_normal_(layer.weight) |
| return layer |
|
|
|
|
| def ConvNd(dims: int = 1, *args, **kwargs): |
| r"""Wrapper of N-dimension convolution with kaiming_normal_ initialization. |
| |
| Args: |
| dims: number of dimensions of the convolution. |
| """ |
| if dims == 1: |
| return Conv1d(*args, **kwargs) |
| elif dims == 2: |
| return Conv2d(*args, **kwargs) |
| else: |
| raise ValueError(f"invalid number of dimensions: {dims}") |
|
|
|
|
| def zero_module(module: nn.Module): |
| r"""Zero out the parameters of a module and return it.""" |
| nn.init.zeros_(module.weight) |
| nn.init.zeros_(module.bias) |
| return module |
|
|
|
|
| def scale_module(module: nn.Module, scale): |
| r"""Scale the parameters of a module and return it.""" |
| for p in module.parameters(): |
| p.detach().mul_(scale) |
| return module |
|
|
|
|
| def mean_flat(tensor: torch.Tensor): |
| r"""Take the mean over all non-batch dimensions.""" |
| return tensor.mean(dim=tuple(range(1, tensor.dim()))) |
|
|
|
|
| def append_dims(x, target_dims): |
| r"""Appends dimensions to the end of a tensor until |
| it has target_dims dimensions. |
| """ |
| dims_to_append = target_dims - x.dim() |
| if dims_to_append < 0: |
| raise ValueError( |
| f"input has {x.dim()} dims but target_dims is {target_dims}, which is less" |
| ) |
| return x[(...,) + (None,) * dims_to_append] |
|
|
|
|
| def append_zero(x, count=1): |
| r"""Appends ``count`` zeros to the end of a tensor along the last dimension.""" |
| assert count > 0, f"invalid count: {count}" |
| return torch.cat([x, x.new_zeros((*x.size()[:-1], count))], dim=-1) |
|
|
|
|
| class Transpose(nn.Identity): |
| """(N, T, D) -> (N, D, T)""" |
|
|
| def forward(self, input: torch.Tensor) -> torch.Tensor: |
| return input.transpose(1, 2) |
|
|
