200 lines
6.6 KiB
Python
200 lines
6.6 KiB
Python
import copy
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
|
|
|
|
class AdaptiveLayerNorm1D(torch.nn.Module):
|
|
def __init__(self, data_dim: int, norm_cond_dim: int):
|
|
super().__init__()
|
|
if data_dim <= 0:
|
|
raise ValueError(f"data_dim must be positive, but got {data_dim}")
|
|
if norm_cond_dim <= 0:
|
|
raise ValueError(f"norm_cond_dim must be positive, but got {norm_cond_dim}")
|
|
self.norm = torch.nn.LayerNorm(
|
|
data_dim
|
|
) # TODO: Check if elementwise_affine=True is correct
|
|
self.linear = torch.nn.Linear(norm_cond_dim, 2 * data_dim)
|
|
torch.nn.init.zeros_(self.linear.weight)
|
|
torch.nn.init.zeros_(self.linear.bias)
|
|
|
|
def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
|
|
# x: (batch, ..., data_dim)
|
|
# t: (batch, norm_cond_dim)
|
|
# return: (batch, data_dim)
|
|
x = self.norm(x)
|
|
alpha, beta = self.linear(t).chunk(2, dim=-1)
|
|
|
|
# Add singleton dimensions to alpha and beta
|
|
if x.dim() > 2:
|
|
alpha = alpha.view(alpha.shape[0], *([1] * (x.dim() - 2)), alpha.shape[1])
|
|
beta = beta.view(beta.shape[0], *([1] * (x.dim() - 2)), beta.shape[1])
|
|
|
|
return x * (1 + alpha) + beta
|
|
|
|
|
|
class SequentialCond(torch.nn.Sequential):
|
|
def forward(self, input, *args, **kwargs):
|
|
for module in self:
|
|
if isinstance(module, (AdaptiveLayerNorm1D, SequentialCond, ResidualMLPBlock)):
|
|
# print(f'Passing on args to {module}', [a.shape for a in args])
|
|
input = module(input, *args, **kwargs)
|
|
else:
|
|
# print(f'Skipping passing args to {module}', [a.shape for a in args])
|
|
input = module(input)
|
|
return input
|
|
|
|
|
|
def normalization_layer(norm: Optional[str], dim: int, norm_cond_dim: int = -1):
|
|
if norm == "batch":
|
|
return torch.nn.BatchNorm1d(dim)
|
|
elif norm == "layer":
|
|
return torch.nn.LayerNorm(dim)
|
|
elif norm == "ada":
|
|
assert norm_cond_dim > 0, f"norm_cond_dim must be positive, got {norm_cond_dim}"
|
|
return AdaptiveLayerNorm1D(dim, norm_cond_dim)
|
|
elif norm is None:
|
|
return torch.nn.Identity()
|
|
else:
|
|
raise ValueError(f"Unknown norm: {norm}")
|
|
|
|
|
|
def linear_norm_activ_dropout(
|
|
input_dim: int,
|
|
output_dim: int,
|
|
activation: torch.nn.Module = torch.nn.ReLU(),
|
|
bias: bool = True,
|
|
norm: Optional[str] = "layer", # Options: ada/batch/layer
|
|
dropout: float = 0.0,
|
|
norm_cond_dim: int = -1,
|
|
) -> SequentialCond:
|
|
layers = []
|
|
layers.append(torch.nn.Linear(input_dim, output_dim, bias=bias))
|
|
if norm is not None:
|
|
layers.append(normalization_layer(norm, output_dim, norm_cond_dim))
|
|
layers.append(copy.deepcopy(activation))
|
|
if dropout > 0.0:
|
|
layers.append(torch.nn.Dropout(dropout))
|
|
return SequentialCond(*layers)
|
|
|
|
|
|
def create_simple_mlp(
|
|
input_dim: int,
|
|
hidden_dims: List[int],
|
|
output_dim: int,
|
|
activation: torch.nn.Module = torch.nn.ReLU(),
|
|
bias: bool = True,
|
|
norm: Optional[str] = "layer", # Options: ada/batch/layer
|
|
dropout: float = 0.0,
|
|
norm_cond_dim: int = -1,
|
|
) -> SequentialCond:
|
|
layers = []
|
|
prev_dim = input_dim
|
|
for hidden_dim in hidden_dims:
|
|
layers.extend(
|
|
linear_norm_activ_dropout(
|
|
prev_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim
|
|
)
|
|
)
|
|
prev_dim = hidden_dim
|
|
layers.append(torch.nn.Linear(prev_dim, output_dim, bias=bias))
|
|
return SequentialCond(*layers)
|
|
|
|
|
|
class ResidualMLPBlock(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
hidden_dim: int,
|
|
num_hidden_layers: int,
|
|
output_dim: int,
|
|
activation: torch.nn.Module = torch.nn.ReLU(),
|
|
bias: bool = True,
|
|
norm: Optional[str] = "layer", # Options: ada/batch/layer
|
|
dropout: float = 0.0,
|
|
norm_cond_dim: int = -1,
|
|
):
|
|
super().__init__()
|
|
if not (input_dim == output_dim == hidden_dim):
|
|
raise NotImplementedError(
|
|
f"input_dim {input_dim} != output_dim {output_dim} is not implemented"
|
|
)
|
|
|
|
layers = []
|
|
prev_dim = input_dim
|
|
for i in range(num_hidden_layers):
|
|
layers.append(
|
|
linear_norm_activ_dropout(
|
|
prev_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim
|
|
)
|
|
)
|
|
prev_dim = hidden_dim
|
|
self.model = SequentialCond(*layers)
|
|
self.skip = torch.nn.Identity()
|
|
|
|
def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
|
|
return x + self.model(x, *args, **kwargs)
|
|
|
|
|
|
class ResidualMLP(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
hidden_dim: int,
|
|
num_hidden_layers: int,
|
|
output_dim: int,
|
|
activation: torch.nn.Module = torch.nn.ReLU(),
|
|
bias: bool = True,
|
|
norm: Optional[str] = "layer", # Options: ada/batch/layer
|
|
dropout: float = 0.0,
|
|
num_blocks: int = 1,
|
|
norm_cond_dim: int = -1,
|
|
):
|
|
super().__init__()
|
|
self.input_dim = input_dim
|
|
self.model = SequentialCond(
|
|
linear_norm_activ_dropout(
|
|
input_dim, hidden_dim, activation, bias, norm, dropout, norm_cond_dim
|
|
),
|
|
*[
|
|
ResidualMLPBlock(
|
|
hidden_dim,
|
|
hidden_dim,
|
|
num_hidden_layers,
|
|
hidden_dim,
|
|
activation,
|
|
bias,
|
|
norm,
|
|
dropout,
|
|
norm_cond_dim,
|
|
)
|
|
for _ in range(num_blocks)
|
|
],
|
|
torch.nn.Linear(hidden_dim, output_dim, bias=bias),
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
|
|
return self.model(x, *args, **kwargs)
|
|
|
|
|
|
class FrequencyEmbedder(torch.nn.Module):
|
|
def __init__(self, num_frequencies, max_freq_log2):
|
|
super().__init__()
|
|
frequencies = 2 ** torch.linspace(0, max_freq_log2, steps=num_frequencies)
|
|
self.register_buffer("frequencies", frequencies)
|
|
|
|
def forward(self, x):
|
|
# x should be of size (N,) or (N, D)
|
|
N = x.size(0)
|
|
if x.dim() == 1: # (N,)
|
|
x = x.unsqueeze(1) # (N, D) where D=1
|
|
x_unsqueezed = x.unsqueeze(-1) # (N, D, 1)
|
|
scaled = self.frequencies.view(1, 1, -1) * x_unsqueezed # (N, D, num_frequencies)
|
|
s = torch.sin(scaled)
|
|
c = torch.cos(scaled)
|
|
embedded = torch.cat([s, c, x_unsqueezed], dim=-1).view(
|
|
N, -1
|
|
) # (N, D * 2 * num_frequencies + D)
|
|
return embedded
|
|
|