File size: 5,761 Bytes
95456ed | 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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | """
Various utilities for neural networks.
"""
import math
import torch
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 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):
"""
Make a standard normalization layer.
:param channels: number of input channels.
:return: an nn.Module for normalization.
"""
return GroupNorm32(32, channels)
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 = 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)
return embedding
def concatenate_sn_sp(sn_repr_emb, sp_repr_emb, sn_repr_len):
# Create an empty tensor to store the concatenated embeddings
concat_emb = torch.empty_like(sn_repr_emb)
# Iterate over the batch size
for i in range(sn_repr_emb.size(0)):
# Get the true sequence length for the sn embedding
seq_len = sn_repr_len[i]
# get the true sequence length for the sp embedding
sp_idx = sp_repr_emb.size(1) - seq_len
# Cut off the extra dimensions in sn_repr_emb and sp_repr_emb
sn_emb = sn_repr_emb[i, :seq_len]
sp_emb = sp_repr_emb[i, :sp_idx]
# Concatenate the embeddings along the sequence length dimension
concat_emb[i] = torch.cat([sn_emb, sp_emb], dim=0)
return concat_emb
def split_into_sn_and_sp(model_output, sn_repr_len):
# create an empty tensor to store the split sn_repr
model_output_sn = torch.empty_like(model_output)
model_output_sp = torch.empty_like(model_output)
microbatch_size, seq_len, emb_dim = model_output.shape
model_output_sn_mask = torch.empty((microbatch_size, seq_len))
model_output_sp_mask = torch.empty((microbatch_size, seq_len))
# iterate over the batc size
for i in range(microbatch_size):
# the start index of the sp output is the length of the sn
orig_sn_len = sn_repr_len[i]
# the SP
# split the sp_repr of the current instance off from the model output
sp_repr_out = model_output[i, orig_sn_len:]
# TODO change this to more sensible padding
# pad the sp representation model output until it is again args.seq_len long
sp_padding = sp_repr_out[-1].repeat(orig_sn_len, 1)
# concatenate the embeddings along the
model_output_sp[i] = torch.cat([sp_repr_out, sp_padding], dim=0)
# the SN
# split the sn_repr of the current instance off from the model output
sn_repr_out = model_output[i, :orig_sn_len]
sn_padding = sn_repr_out[-1].repeat(seq_len - orig_sn_len, 1)
model_output_sn[i] = torch.cat([sn_repr_out, sn_padding], dim=0)
# the masks: they only mask the additional padding added now, not the original padding added after the sp
sp_mask = torch.ones(seq_len).to(model_output.device)
sp_mask[-orig_sn_len:] = 0
model_output_sp_mask[i] = sp_mask
sn_mask = torch.ones(seq_len).to(model_output.device)
sn_mask[orig_sn_len:] = 0
model_output_sn_mask[i] = sn_mask
model_output_sn_mask = model_output_sn_mask.to(model_output.device)
model_output_sp_mask = model_output_sp_mask.to(model_output.device)
return model_output_sn, model_output_sp, model_output_sn_mask, model_output_sp_mask
|