Spaces:
Sleeping
Sleeping
File size: 3,745 Bytes
5007d4b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# This code is based on https://github.com/openai/guided-diffusion
"""
Various utilities for neural networks.
"""
import torch as th
import torch.nn as nn
# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
class SiLU(nn.Module):
def forward(self, x):
return x * th.sigmoid(x)
class GroupNorm32(nn.GroupNorm):
def forward(self, x):
return super().forward(x.float()).type(x.dtype)
def linear(*args, **kwargs):
"""
Create a linear module.
"""
return nn.Linear(*args, **kwargs)
def mean_flat(tensor):
"""
Take the mean over all non-batch dimensions.
"""
return tensor.mean(dim=list(range(1, len(tensor.shape))))
def sum_flat(tensor):
"""
Take the sum over all non-batch dimensions.
"""
return tensor.sum(dim=list(range(1, len(tensor.shape))))
def normalization(channels):
"""
Make a standard normalization layer.
:param channels: number of input channels.
:return: an nn.Module for normalization.
"""
return GroupNorm32(32, channels)
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
@th.amp.custom_fwd(device_type="cuda")
def forward(ctx, run_function, length, *args):
ctx.run_function = run_function
ctx.input_length = length
ctx.save_for_backward(*args)
with th.no_grad():
output_tensors = ctx.run_function(*args[:length])
return output_tensors
@staticmethod
@th.amp.custom_bwd(device_type="cuda")
def backward(ctx, *output_grads):
args = list(ctx.saved_tensors)
# Filter for inputs that require grad. If none, exit early.
input_indices = [i for (i, x) in enumerate(args) if x.requires_grad]
if not input_indices:
return (None, None) + tuple(None for _ in args)
with th.enable_grad():
for i in input_indices:
if i < ctx.input_length:
# Not sure why the OAI code does this little
# dance. It might not be necessary.
args[i] = args[i].detach().requires_grad_()
args[i] = args[i].view_as(args[i])
output_tensors = ctx.run_function(*args[: ctx.input_length])
if isinstance(output_tensors, th.Tensor):
output_tensors = [output_tensors]
# Filter for outputs that require grad. If none, exit early.
out_and_grads = [
(o, g) for (o, g) in zip(output_tensors, output_grads) if o.requires_grad
]
if not out_and_grads:
return (None, None) + tuple(None for _ in args)
# Compute gradients on the filtered tensors.
computed_grads = th.autograd.grad(
[o for (o, g) in out_and_grads],
[args[i] for i in input_indices],
[g for (o, g) in out_and_grads],
)
# Reassemble the complete gradient tuple.
input_grads = [None for _ in args]
for i, g in zip(input_indices, computed_grads):
input_grads[i] = g
return (None, None) + tuple(input_grads)
|