|
import math |
|
import torch |
|
import torch.nn as nn |
|
from einops import repeat |
|
|
|
|
|
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(torch.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 torch.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 torch.enable_grad(): |
|
|
|
|
|
|
|
shallow_copies = [x.view_as(x) for x in ctx.input_tensors] |
|
output_tensors = ctx.run_function(*shallow_copies) |
|
input_grads = torch.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 |
|
|
|
|
|
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 = torch.exp( |
|
-math.log(max_period) |
|
* torch.arange(start=0, end=half, dtype=torch.float32) |
|
/ half |
|
).to(device=timesteps.device) |
|
args = timesteps[:, None] * 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 |
|
) |
|
else: |
|
embedding = repeat(timesteps, "b -> b d", d=dim) |
|
|
|
return embedding |
|
|
|
|
|
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 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 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}") |
|
|
