| |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| import torch |
| import torch.nn as nn |
| import einops |
|
|
| from inspect import isfunction |
|
|
|
|
| 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 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 nonlinearity(type='silu'): |
| if type == 'silu': |
| return nn.SiLU() |
| elif type == 'leaky_relu': |
| return nn.LeakyReLU() |
|
|
|
|
| def normalization(channels, num_groups=32): |
| """ |
| Make a standard normalization layer. |
| :param channels: number of input channels. |
| :return: an nn.Module for normalization. |
| """ |
| return nn.GroupNorm(num_groups, channels) |
|
|
|
|
| def default(val, d): |
| if exists(val): |
| return val |
| return d() if isfunction(d) else d |
|
|
|
|
| def exists(val): |
| return val is not None |
|
|
|
|
| def extract_into_tensor(a, t, x_shape): |
| b, *_ = t.shape |
| out = a.gather(-1, t) |
| return out.reshape(b, *((1,) * (len(x_shape) - 1))) |
|
|
|
|
| def make_temporal_window(x, t, method): |
| assert method in ['roll', 'prv', 'first'] |
|
|
| if method == 'roll': |
| m = einops.rearrange(x, '(b t) d c -> b t d c', t=t) |
| l = torch.roll(m, shifts=1, dims=1) |
| r = torch.roll(m, shifts=-1, dims=1) |
|
|
| recon = torch.cat([l, m, r], dim=2) |
| del l, m, r |
|
|
| recon = einops.rearrange(recon, 'b t d c -> (b t) d c') |
| return recon |
|
|
| if method == 'prv': |
| x = einops.rearrange(x, '(b t) d c -> b t d c', t=t) |
| prv = torch.cat([x[:, :1], x[:, :-1]], dim=1) |
|
|
| recon = torch.cat([x, prv], dim=2) |
| del x, prv |
|
|
| recon = einops.rearrange(recon, 'b t d c -> (b t) d c') |
| return recon |
|
|
| if method == 'first': |
| x = einops.rearrange(x, '(b t) d c -> b t d c', t=t) |
| prv = x[:, [0], :, :].repeat(1, t, 1, 1) |
|
|
| recon = torch.cat([x, prv], dim=2) |
| del x, prv |
|
|
| recon = einops.rearrange(recon, 'b t d c -> (b t) d c') |
| return recon |
|
|
|
|
| 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: |
| return torch.utils.checkpoint.checkpoint(func, *inputs, use_reentrant=False) |
| else: |
| return func(*inputs) |
|
|
