|
""" |
|
Various utilities for neural networks. |
|
""" |
|
|
|
import math |
|
|
|
import torch as th |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
|
|
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 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) |
|
|
|
|
|
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 mean_flat(tensor): |
|
""" |
|
Take the mean over all non-batch dimensions. |
|
""" |
|
return tensor.mean(dim=list(range(1, len(tensor.shape)))) |
|
|
|
|
|
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, repeat_only=False): |
|
""" |
|
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. |
|
""" |
|
if not repeat_only: |
|
half = dim // 2 |
|
freqs = th.exp( |
|
-math.log(max_period) |
|
* th.arange(start=0, end=half, dtype=th.float32) |
|
/ half |
|
).to(device=timesteps.device) |
|
args = timesteps[:, None].float() * freqs[None] |
|
embedding = th.cat([th.cos(args), th.sin(args)], dim=-1) |
|
if dim % 2: |
|
embedding = th.cat([embedding, th.zeros_like(embedding[:, :1])], dim=-1) |
|
else: |
|
embedding = repeat(timesteps, "b -> b d", d=dim) |
|
return embedding |
|
|
|
|
|
def checkpoint(func, inputs, params, flag): |
|
""" |
|
Evaluate a function without caching intermediate activations, allowing for |
|
reduced memory at the expense of extra compute in the backward pass. |
|
|
|
:param func: the function to evaluate. |
|
:param inputs: the argument sequence to pass to `func`. |
|
:param params: a sequence of parameters `func` depends on but does not |
|
explicitly take as arguments. |
|
:param flag: if False, disable gradient checkpointing. |
|
""" |
|
|
|
if flag: |
|
args = tuple(inputs) + tuple(params) |
|
return CheckpointFunction.apply(func, len(inputs), *args) |
|
else: |
|
return func(*inputs) |
|
|
|
|
|
class CheckpointFunction(th.autograd.Function): |
|
@staticmethod |
|
def forward(ctx, run_function, length, *args): |
|
ctx.run_function = run_function |
|
ctx.input_tensors = list(args[:length]) |
|
ctx.input_params = list(args[length:]) |
|
with th.no_grad(): |
|
output_tensors = ctx.run_function(*ctx.input_tensors) |
|
return output_tensors |
|
|
|
@staticmethod |
|
def backward(ctx, *output_grads): |
|
ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors] |
|
with th.enable_grad(): |
|
|
|
|
|
|
|
shallow_copies = [x.view_as(x) for x in ctx.input_tensors] |
|
output_tensors = ctx.run_function(*shallow_copies) |
|
input_grads = th.autograd.grad( |
|
output_tensors, |
|
ctx.input_tensors + ctx.input_params, |
|
output_grads, |
|
allow_unused=True, |
|
) |
|
del ctx.input_tensors |
|
del ctx.input_params |
|
del output_tensors |
|
return (None, None) + input_grads |
|
|
