359 lines
11 KiB
Python
359 lines
11 KiB
Python
from inspect import isfunction
|
|
from typing import Callable, Optional
|
|
|
|
import torch
|
|
from einops import rearrange
|
|
from einops.layers.torch import Rearrange
|
|
from torch import nn
|
|
|
|
from .t_cond_mlp import (
|
|
AdaptiveLayerNorm1D,
|
|
FrequencyEmbedder,
|
|
normalization_layer,
|
|
)
|
|
# from .vit import Attention, FeedForward
|
|
|
|
|
|
def exists(val):
|
|
return val is not None
|
|
|
|
|
|
def default(val, d):
|
|
if exists(val):
|
|
return val
|
|
return d() if isfunction(d) else d
|
|
|
|
|
|
class PreNorm(nn.Module):
|
|
def __init__(self, dim: int, fn: Callable, norm: str = "layer", norm_cond_dim: int = -1):
|
|
super().__init__()
|
|
self.norm = normalization_layer(norm, dim, norm_cond_dim)
|
|
self.fn = fn
|
|
|
|
def forward(self, x: torch.Tensor, *args, **kwargs):
|
|
if isinstance(self.norm, AdaptiveLayerNorm1D):
|
|
return self.fn(self.norm(x, *args), **kwargs)
|
|
else:
|
|
return self.fn(self.norm(x), **kwargs)
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, dim, hidden_dim, dropout=0.0):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Linear(dim, hidden_dim),
|
|
nn.GELU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(hidden_dim, dim),
|
|
nn.Dropout(dropout),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(self, dim, heads=8, dim_head=64, dropout=0.0):
|
|
super().__init__()
|
|
inner_dim = dim_head * heads
|
|
project_out = not (heads == 1 and dim_head == dim)
|
|
|
|
self.heads = heads
|
|
self.scale = dim_head**-0.5
|
|
|
|
self.attend = nn.Softmax(dim=-1)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
|
|
|
|
self.to_out = (
|
|
nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
|
|
if project_out
|
|
else nn.Identity()
|
|
)
|
|
|
|
def forward(self, x):
|
|
qkv = self.to_qkv(x).chunk(3, dim=-1)
|
|
q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), qkv)
|
|
|
|
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
|
|
|
|
attn = self.attend(dots)
|
|
attn = self.dropout(attn)
|
|
|
|
out = torch.matmul(attn, v)
|
|
out = rearrange(out, "b h n d -> b n (h d)")
|
|
return self.to_out(out)
|
|
|
|
|
|
class CrossAttention(nn.Module):
|
|
def __init__(self, dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
|
|
super().__init__()
|
|
inner_dim = dim_head * heads
|
|
project_out = not (heads == 1 and dim_head == dim)
|
|
|
|
self.heads = heads
|
|
self.scale = dim_head**-0.5
|
|
|
|
self.attend = nn.Softmax(dim=-1)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
context_dim = default(context_dim, dim)
|
|
self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=False)
|
|
self.to_q = nn.Linear(dim, inner_dim, bias=False)
|
|
|
|
self.to_out = (
|
|
nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
|
|
if project_out
|
|
else nn.Identity()
|
|
)
|
|
|
|
def forward(self, x, context=None):
|
|
context = default(context, x)
|
|
k, v = self.to_kv(context).chunk(2, dim=-1)
|
|
q = self.to_q(x)
|
|
q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), [q, k, v])
|
|
|
|
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
|
|
|
|
attn = self.attend(dots)
|
|
attn = self.dropout(attn)
|
|
|
|
out = torch.matmul(attn, v)
|
|
out = rearrange(out, "b h n d -> b n (h d)")
|
|
return self.to_out(out)
|
|
|
|
|
|
class Transformer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
depth: int,
|
|
heads: int,
|
|
dim_head: int,
|
|
mlp_dim: int,
|
|
dropout: float = 0.0,
|
|
norm: str = "layer",
|
|
norm_cond_dim: int = -1,
|
|
):
|
|
super().__init__()
|
|
self.layers = nn.ModuleList([])
|
|
for _ in range(depth):
|
|
sa = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)
|
|
ff = FeedForward(dim, mlp_dim, dropout=dropout)
|
|
self.layers.append(
|
|
nn.ModuleList(
|
|
[
|
|
PreNorm(dim, sa, norm=norm, norm_cond_dim=norm_cond_dim),
|
|
PreNorm(dim, ff, norm=norm, norm_cond_dim=norm_cond_dim),
|
|
]
|
|
)
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, *args):
|
|
for attn, ff in self.layers:
|
|
x = attn(x, *args) + x
|
|
x = ff(x, *args) + x
|
|
return x
|
|
|
|
|
|
class TransformerCrossAttn(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
depth: int,
|
|
heads: int,
|
|
dim_head: int,
|
|
mlp_dim: int,
|
|
dropout: float = 0.0,
|
|
norm: str = "layer",
|
|
norm_cond_dim: int = -1,
|
|
context_dim: Optional[int] = None,
|
|
):
|
|
super().__init__()
|
|
self.layers = nn.ModuleList([])
|
|
for _ in range(depth):
|
|
sa = Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)
|
|
ca = CrossAttention(
|
|
dim, context_dim=context_dim, heads=heads, dim_head=dim_head, dropout=dropout
|
|
)
|
|
ff = FeedForward(dim, mlp_dim, dropout=dropout)
|
|
self.layers.append(
|
|
nn.ModuleList(
|
|
[
|
|
PreNorm(dim, sa, norm=norm, norm_cond_dim=norm_cond_dim),
|
|
PreNorm(dim, ca, norm=norm, norm_cond_dim=norm_cond_dim),
|
|
PreNorm(dim, ff, norm=norm, norm_cond_dim=norm_cond_dim),
|
|
]
|
|
)
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, *args, context=None, context_list=None):
|
|
if context_list is None:
|
|
context_list = [context] * len(self.layers)
|
|
if len(context_list) != len(self.layers):
|
|
raise ValueError(f"len(context_list) != len(self.layers) ({len(context_list)} != {len(self.layers)})")
|
|
|
|
for i, (self_attn, cross_attn, ff) in enumerate(self.layers):
|
|
x = self_attn(x, *args) + x
|
|
x = cross_attn(x, *args, context=context_list[i]) + x
|
|
x = ff(x, *args) + x
|
|
return x
|
|
|
|
|
|
class DropTokenDropout(nn.Module):
|
|
def __init__(self, p: float = 0.1):
|
|
super().__init__()
|
|
if p < 0 or p > 1:
|
|
raise ValueError(
|
|
"dropout probability has to be between 0 and 1, " "but got {}".format(p)
|
|
)
|
|
self.p = p
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
# x: (batch_size, seq_len, dim)
|
|
if self.training and self.p > 0:
|
|
zero_mask = torch.full_like(x[0, :, 0], self.p).bernoulli().bool()
|
|
# TODO: permutation idx for each batch using torch.argsort
|
|
if zero_mask.any():
|
|
x = x[:, ~zero_mask, :]
|
|
return x
|
|
|
|
|
|
class ZeroTokenDropout(nn.Module):
|
|
def __init__(self, p: float = 0.1):
|
|
super().__init__()
|
|
if p < 0 or p > 1:
|
|
raise ValueError(
|
|
"dropout probability has to be between 0 and 1, " "but got {}".format(p)
|
|
)
|
|
self.p = p
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
# x: (batch_size, seq_len, dim)
|
|
if self.training and self.p > 0:
|
|
zero_mask = torch.full_like(x[:, :, 0], self.p).bernoulli().bool()
|
|
# Zero-out the masked tokens
|
|
x[zero_mask, :] = 0
|
|
return x
|
|
|
|
|
|
class TransformerEncoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_tokens: int,
|
|
token_dim: int,
|
|
dim: int,
|
|
depth: int,
|
|
heads: int,
|
|
mlp_dim: int,
|
|
dim_head: int = 64,
|
|
dropout: float = 0.0,
|
|
emb_dropout: float = 0.0,
|
|
emb_dropout_type: str = "drop",
|
|
emb_dropout_loc: str = "token",
|
|
norm: str = "layer",
|
|
norm_cond_dim: int = -1,
|
|
token_pe_numfreq: int = -1,
|
|
):
|
|
super().__init__()
|
|
if token_pe_numfreq > 0:
|
|
token_dim_new = token_dim * (2 * token_pe_numfreq + 1)
|
|
self.to_token_embedding = nn.Sequential(
|
|
Rearrange("b n d -> (b n) d", n=num_tokens, d=token_dim),
|
|
FrequencyEmbedder(token_pe_numfreq, token_pe_numfreq - 1),
|
|
Rearrange("(b n) d -> b n d", n=num_tokens, d=token_dim_new),
|
|
nn.Linear(token_dim_new, dim),
|
|
)
|
|
else:
|
|
self.to_token_embedding = nn.Linear(token_dim, dim)
|
|
self.pos_embedding = nn.Parameter(torch.randn(1, num_tokens, dim))
|
|
if emb_dropout_type == "drop":
|
|
self.dropout = DropTokenDropout(emb_dropout)
|
|
elif emb_dropout_type == "zero":
|
|
self.dropout = ZeroTokenDropout(emb_dropout)
|
|
else:
|
|
raise ValueError(f"Unknown emb_dropout_type: {emb_dropout_type}")
|
|
self.emb_dropout_loc = emb_dropout_loc
|
|
|
|
self.transformer = Transformer(
|
|
dim, depth, heads, dim_head, mlp_dim, dropout, norm=norm, norm_cond_dim=norm_cond_dim
|
|
)
|
|
|
|
def forward(self, inp: torch.Tensor, *args, **kwargs):
|
|
x = inp
|
|
|
|
if self.emb_dropout_loc == "input":
|
|
x = self.dropout(x)
|
|
x = self.to_token_embedding(x)
|
|
|
|
if self.emb_dropout_loc == "token":
|
|
x = self.dropout(x)
|
|
b, n, _ = x.shape
|
|
x += self.pos_embedding[:, :n]
|
|
|
|
if self.emb_dropout_loc == "token_afterpos":
|
|
x = self.dropout(x)
|
|
x = self.transformer(x, *args)
|
|
return x
|
|
|
|
|
|
class TransformerDecoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
num_tokens: int,
|
|
token_dim: int,
|
|
dim: int,
|
|
depth: int,
|
|
heads: int,
|
|
mlp_dim: int,
|
|
dim_head: int = 64,
|
|
dropout: float = 0.0,
|
|
emb_dropout: float = 0.0,
|
|
emb_dropout_type: str = 'drop',
|
|
norm: str = "layer",
|
|
norm_cond_dim: int = -1,
|
|
context_dim: Optional[int] = None,
|
|
skip_token_embedding: bool = False,
|
|
):
|
|
super().__init__()
|
|
if not skip_token_embedding:
|
|
self.to_token_embedding = nn.Linear(token_dim, dim)
|
|
else:
|
|
self.to_token_embedding = nn.Identity()
|
|
if token_dim != dim:
|
|
raise ValueError(
|
|
f"token_dim ({token_dim}) != dim ({dim}) when skip_token_embedding is True"
|
|
)
|
|
|
|
self.pos_embedding = nn.Parameter(torch.randn(1, num_tokens, dim))
|
|
if emb_dropout_type == "drop":
|
|
self.dropout = DropTokenDropout(emb_dropout)
|
|
elif emb_dropout_type == "zero":
|
|
self.dropout = ZeroTokenDropout(emb_dropout)
|
|
elif emb_dropout_type == "normal":
|
|
self.dropout = nn.Dropout(emb_dropout)
|
|
|
|
self.transformer = TransformerCrossAttn(
|
|
dim,
|
|
depth,
|
|
heads,
|
|
dim_head,
|
|
mlp_dim,
|
|
dropout,
|
|
norm=norm,
|
|
norm_cond_dim=norm_cond_dim,
|
|
context_dim=context_dim,
|
|
)
|
|
|
|
def forward(self, inp: torch.Tensor, *args, context=None, context_list=None):
|
|
x = self.to_token_embedding(inp)
|
|
b, n, _ = x.shape
|
|
|
|
x = self.dropout(x)
|
|
x += self.pos_embedding[:, :n]
|
|
|
|
x = self.transformer(x, *args, context=context, context_list=context_list)
|
|
return x
|
|
|