122 lines
3.4 KiB
Python
122 lines
3.4 KiB
Python
"""
|
|
Various utilities for neural networks.
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def update_ema(target_params, source_params, rate=0.99):
|
|
"""
|
|
Update target parameters to be closer to those of source parameters using
|
|
an exponential moving average.
|
|
:param target_params: the target parameter sequence.
|
|
:param source_params: the source parameter sequence.
|
|
:param rate: the EMA rate (closer to 1 means slower).
|
|
"""
|
|
for targ, src in zip(target_params, source_params):
|
|
targ.detach().mul_(rate).add_(src, alpha=1 - rate)
|
|
|
|
return target_params
|
|
|
|
|
|
class GroupNorm32(nn.GroupNorm):
|
|
def __init__(self, num_groups, num_channels, swish, eps=1e-5):
|
|
super().__init__(num_groups=num_groups, num_channels=num_channels, eps=eps)
|
|
self.swish = swish
|
|
|
|
def forward(self, x):
|
|
y = super().forward(x.float()).to(x.dtype)
|
|
if self.swish == 1.0:
|
|
y = F.silu(y)
|
|
elif self.swish:
|
|
y = y * F.sigmoid(y * float(self.swish))
|
|
return y
|
|
|
|
|
|
def conv_nd(dims, *args, **kwargs):
|
|
"""
|
|
Create a 1D, 2D, or 3D convolution module.
|
|
"""
|
|
if dims == 1:
|
|
return nn.Conv1d(*args, **kwargs)
|
|
elif dims == 2:
|
|
return nn.Conv2d(*args, **kwargs)
|
|
elif dims == 3:
|
|
return nn.Conv3d(*args, **kwargs)
|
|
raise ValueError(f"unsupported dimensions: {dims}")
|
|
|
|
|
|
def linear(*args, **kwargs):
|
|
"""
|
|
Create a linear module.
|
|
"""
|
|
return nn.Linear(*args, **kwargs)
|
|
|
|
|
|
def avg_pool_nd(dims, *args, **kwargs):
|
|
"""
|
|
Create a 1D, 2D, or 3D average pooling module.
|
|
"""
|
|
if dims == 1:
|
|
return nn.AvgPool1d(*args, **kwargs)
|
|
elif dims == 2:
|
|
return nn.AvgPool2d(*args, **kwargs)
|
|
elif dims == 3:
|
|
return nn.AvgPool3d(*args, **kwargs)
|
|
raise ValueError(f"unsupported dimensions: {dims}")
|
|
|
|
|
|
def zero_module(module):
|
|
"""
|
|
Zero out the parameters of a module and return it.
|
|
"""
|
|
for p in module.parameters():
|
|
p.detach().zero_()
|
|
return module
|
|
|
|
|
|
def scale_module(module, scale):
|
|
"""
|
|
Scale the parameters of a module and return it.
|
|
"""
|
|
for p in module.parameters():
|
|
p.detach().mul_(scale)
|
|
return module
|
|
|
|
|
|
def normalization(channels, swish=0.0):
|
|
"""
|
|
Make a standard normalization layer, with an optional swish activation.
|
|
|
|
:param channels: number of input channels.
|
|
:return: an nn.Module for normalization.
|
|
"""
|
|
return GroupNorm32(num_channels=channels, num_groups=32, swish=swish)
|
|
|
|
|
|
def timestep_embedding(timesteps, dim, max_period=10000):
|
|
"""
|
|
Create sinusoidal timestep embeddings.
|
|
|
|
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
|
These may be fractional.
|
|
:param dim: the dimension of the output.
|
|
:param max_period: controls the minimum frequency of the embeddings.
|
|
:return: an [N x dim] Tensor of positional embeddings.
|
|
"""
|
|
half = dim // 2
|
|
freqs = torch.exp(
|
|
-math.log(max_period)
|
|
* torch.arange(start=0, end=half, dtype=torch.float32)
|
|
/ half
|
|
).to(device=timesteps.device)
|
|
args = timesteps[:, None].float() * freqs[None]
|
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
if dim % 2:
|
|
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
|
return embedding
|