mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
730 lines
21 KiB
Python
730 lines
21 KiB
Python
"""shout-out to https://github.com/lucidrains/x-transformers/tree/main/x_transformers"""
|
|
from collections import namedtuple
|
|
from functools import partial
|
|
from inspect import isfunction
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from einops import rearrange, reduce, repeat
|
|
from torch import einsum, nn
|
|
|
|
# constants
|
|
|
|
DEFAULT_DIM_HEAD = 64
|
|
|
|
Intermediates = namedtuple("Intermediates", ["pre_softmax_attn", "post_softmax_attn"])
|
|
|
|
LayerIntermediates = namedtuple("Intermediates", ["hiddens", "attn_intermediates"])
|
|
|
|
|
|
class AbsolutePositionalEmbedding(nn.Module):
|
|
def __init__(self, dim, max_seq_len):
|
|
super().__init__()
|
|
self.emb = nn.Embedding(max_seq_len, dim)
|
|
self.init_()
|
|
|
|
def init_(self):
|
|
nn.init.normal_(self.emb.weight, std=0.02)
|
|
|
|
def forward(self, x):
|
|
n = torch.arange(x.shape[1], device=x.device)
|
|
return self.emb(n)[None, :, :]
|
|
|
|
|
|
class FixedPositionalEmbedding(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
|
|
self.register_buffer("inv_freq", inv_freq)
|
|
|
|
def forward(self, x, seq_dim=1, offset=0):
|
|
t = (
|
|
torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq)
|
|
+ offset
|
|
)
|
|
sinusoid_inp = torch.einsum("i , j -> i j", t, self.inv_freq)
|
|
emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
|
|
return emb[None, :, :]
|
|
|
|
|
|
# helpers
|
|
|
|
|
|
def exists(val):
|
|
return val is not None
|
|
|
|
|
|
def default(val, d):
|
|
if exists(val):
|
|
return val
|
|
return d() if isfunction(d) else d
|
|
|
|
|
|
def always(val):
|
|
def inner(*args, **kwargs):
|
|
return val
|
|
|
|
return inner
|
|
|
|
|
|
def not_equals(val):
|
|
def inner(x):
|
|
return x != val
|
|
|
|
return inner
|
|
|
|
|
|
def equals(val):
|
|
def inner(x):
|
|
return x == val
|
|
|
|
return inner
|
|
|
|
|
|
def max_neg_value(tensor):
|
|
return -torch.finfo(tensor.dtype).max
|
|
|
|
|
|
# keyword argument helpers
|
|
|
|
|
|
def pick_and_pop(keys, d):
|
|
values = list(map(lambda key: d.pop(key), keys))
|
|
return dict(zip(keys, values))
|
|
|
|
|
|
def group_dict_by_key(cond, d):
|
|
return_val = [dict(), dict()]
|
|
for key in d.keys():
|
|
match = bool(cond(key))
|
|
ind = int(not match)
|
|
return_val[ind][key] = d[key]
|
|
return (*return_val,)
|
|
|
|
|
|
def string_begins_with(prefix, str):
|
|
return str.startswith(prefix)
|
|
|
|
|
|
def group_by_key_prefix(prefix, d):
|
|
return group_dict_by_key(partial(string_begins_with, prefix), d)
|
|
|
|
|
|
def groupby_prefix_and_trim(prefix, d):
|
|
kwargs_with_prefix, kwargs = group_dict_by_key(
|
|
partial(string_begins_with, prefix), d
|
|
)
|
|
kwargs_without_prefix = dict(
|
|
map(
|
|
lambda x: (x[0][len(prefix) :], x[1]),
|
|
tuple(kwargs_with_prefix.items()),
|
|
)
|
|
)
|
|
return kwargs_without_prefix, kwargs
|
|
|
|
|
|
# classes
|
|
class Scale(nn.Module):
|
|
def __init__(self, value, fn):
|
|
super().__init__()
|
|
self.value = value
|
|
self.fn = fn
|
|
|
|
def forward(self, x, **kwargs):
|
|
x, *rest = self.fn(x, **kwargs)
|
|
return (x * self.value, *rest)
|
|
|
|
|
|
class Rezero(nn.Module):
|
|
def __init__(self, fn):
|
|
super().__init__()
|
|
self.fn = fn
|
|
self.g = nn.Parameter(torch.zeros(1))
|
|
|
|
def forward(self, x, **kwargs):
|
|
x, *rest = self.fn(x, **kwargs)
|
|
return (x * self.g, *rest)
|
|
|
|
|
|
class ScaleNorm(nn.Module):
|
|
def __init__(self, dim, eps=1e-5):
|
|
super().__init__()
|
|
self.scale = dim**-0.5
|
|
self.eps = eps
|
|
self.g = nn.Parameter(torch.ones(1))
|
|
|
|
def forward(self, x):
|
|
norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
|
|
return x / norm.clamp(min=self.eps) * self.g
|
|
|
|
|
|
class RMSNorm(nn.Module):
|
|
def __init__(self, dim, eps=1e-8):
|
|
super().__init__()
|
|
self.scale = dim**-0.5
|
|
self.eps = eps
|
|
self.g = nn.Parameter(torch.ones(dim))
|
|
|
|
def forward(self, x):
|
|
norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
|
|
return x / norm.clamp(min=self.eps) * self.g
|
|
|
|
|
|
class Residual(nn.Module):
|
|
def forward(self, x, residual):
|
|
return x + residual
|
|
|
|
|
|
class GRUGating(nn.Module):
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.gru = nn.GRUCell(dim, dim)
|
|
|
|
def forward(self, x, residual):
|
|
gated_output = self.gru(
|
|
rearrange(x, "b n d -> (b n) d"),
|
|
rearrange(residual, "b n d -> (b n) d"),
|
|
)
|
|
|
|
return gated_output.reshape_as(x)
|
|
|
|
|
|
# feedforward
|
|
|
|
|
|
class GEGLU(nn.Module):
|
|
def __init__(self, dim_in, dim_out):
|
|
super().__init__()
|
|
self.proj = nn.Linear(dim_in, dim_out * 2)
|
|
|
|
def forward(self, x):
|
|
x, gate = self.proj(x).chunk(2, dim=-1)
|
|
return x * F.gelu(gate)
|
|
|
|
|
|
class FeedForward(nn.Module):
|
|
def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
|
|
super().__init__()
|
|
inner_dim = int(dim * mult)
|
|
dim_out = default(dim_out, dim)
|
|
project_in = (
|
|
nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
|
|
if not glu
|
|
else GEGLU(dim, inner_dim)
|
|
)
|
|
|
|
self.net = nn.Sequential(
|
|
project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
# attention.
|
|
class Attention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
dim_head=DEFAULT_DIM_HEAD,
|
|
heads=8,
|
|
causal=False,
|
|
mask=None,
|
|
talking_heads=False,
|
|
sparse_topk=None,
|
|
use_entmax15=False,
|
|
num_mem_kv=0,
|
|
dropout=0.0,
|
|
on_attn=False,
|
|
):
|
|
super().__init__()
|
|
if use_entmax15:
|
|
raise NotImplementedError(
|
|
"Check out entmax activation instead of softmax activation!"
|
|
)
|
|
self.scale = dim_head**-0.5
|
|
self.heads = heads
|
|
self.causal = causal
|
|
self.mask = mask
|
|
|
|
inner_dim = dim_head * heads
|
|
|
|
self.to_q = nn.Linear(dim, inner_dim, bias=False)
|
|
self.to_k = nn.Linear(dim, inner_dim, bias=False)
|
|
self.to_v = nn.Linear(dim, inner_dim, bias=False)
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
# talking heads
|
|
self.talking_heads = talking_heads
|
|
if talking_heads:
|
|
self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
|
|
self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
|
|
|
|
# explicit topk sparse attention
|
|
self.sparse_topk = sparse_topk
|
|
|
|
# entmax
|
|
# self.attn_fn = entmax15 if use_entmax15 else F.softmax
|
|
self.attn_fn = F.softmax
|
|
|
|
# add memory key / values
|
|
self.num_mem_kv = num_mem_kv
|
|
if num_mem_kv > 0:
|
|
self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
|
|
self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
|
|
|
|
# attention on attention
|
|
self.attn_on_attn = on_attn
|
|
self.to_out = (
|
|
nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU())
|
|
if on_attn
|
|
else nn.Linear(inner_dim, dim)
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
context=None,
|
|
mask=None,
|
|
context_mask=None,
|
|
rel_pos=None,
|
|
sinusoidal_emb=None,
|
|
prev_attn=None,
|
|
mem=None,
|
|
):
|
|
b, n, _, h, talking_heads, device = (
|
|
*x.shape,
|
|
self.heads,
|
|
self.talking_heads,
|
|
x.device,
|
|
)
|
|
kv_input = default(context, x)
|
|
|
|
q_input = x
|
|
k_input = kv_input
|
|
v_input = kv_input
|
|
|
|
if exists(mem):
|
|
k_input = torch.cat((mem, k_input), dim=-2)
|
|
v_input = torch.cat((mem, v_input), dim=-2)
|
|
|
|
if exists(sinusoidal_emb):
|
|
# in shortformer, the query would start at a position offset depending on the past cached memory
|
|
offset = k_input.shape[-2] - q_input.shape[-2]
|
|
q_input = q_input + sinusoidal_emb(q_input, offset=offset)
|
|
k_input = k_input + sinusoidal_emb(k_input)
|
|
|
|
q = self.to_q(q_input)
|
|
k = self.to_k(k_input)
|
|
v = self.to_v(v_input)
|
|
|
|
q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v))
|
|
|
|
input_mask = None
|
|
if any(map(exists, (mask, context_mask))):
|
|
q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
|
|
k_mask = q_mask if not exists(context) else context_mask
|
|
k_mask = default(
|
|
k_mask,
|
|
lambda: torch.ones((b, k.shape[-2]), device=device).bool(),
|
|
)
|
|
q_mask = rearrange(q_mask, "b i -> b () i ()")
|
|
k_mask = rearrange(k_mask, "b j -> b () () j")
|
|
input_mask = q_mask * k_mask
|
|
|
|
if self.num_mem_kv > 0:
|
|
mem_k, mem_v = map(
|
|
lambda t: repeat(t, "h n d -> b h n d", b=b),
|
|
(self.mem_k, self.mem_v),
|
|
)
|
|
k = torch.cat((mem_k, k), dim=-2)
|
|
v = torch.cat((mem_v, v), dim=-2)
|
|
if exists(input_mask):
|
|
input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
|
|
|
|
dots = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
|
|
mask_value = max_neg_value(dots)
|
|
|
|
if exists(prev_attn):
|
|
dots = dots + prev_attn
|
|
|
|
pre_softmax_attn = dots
|
|
|
|
if talking_heads:
|
|
dots = einsum(
|
|
"b h i j, h k -> b k i j", dots, self.pre_softmax_proj
|
|
).contiguous()
|
|
|
|
if exists(rel_pos):
|
|
dots = rel_pos(dots)
|
|
|
|
if exists(input_mask):
|
|
dots.masked_fill_(~input_mask, mask_value)
|
|
del input_mask
|
|
|
|
if self.causal:
|
|
i, j = dots.shape[-2:]
|
|
r = torch.arange(i, device=device)
|
|
mask = rearrange(r, "i -> () () i ()") < rearrange(r, "j -> () () () j")
|
|
mask = F.pad(mask, (j - i, 0), value=False)
|
|
dots.masked_fill_(mask, mask_value)
|
|
del mask
|
|
|
|
if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
|
|
top, _ = dots.topk(self.sparse_topk, dim=-1)
|
|
vk = top[..., -1].unsqueeze(-1).expand_as(dots)
|
|
mask = dots < vk
|
|
dots.masked_fill_(mask, mask_value)
|
|
del mask
|
|
|
|
attn = self.attn_fn(dots, dim=-1)
|
|
post_softmax_attn = attn
|
|
|
|
attn = self.dropout(attn)
|
|
|
|
if talking_heads:
|
|
attn = einsum(
|
|
"b h i j, h k -> b k i j", attn, self.post_softmax_proj
|
|
).contiguous()
|
|
|
|
out = einsum("b h i j, b h j d -> b h i d", attn, v)
|
|
out = rearrange(out, "b h n d -> b n (h d)")
|
|
|
|
intermediates = Intermediates(
|
|
pre_softmax_attn=pre_softmax_attn,
|
|
post_softmax_attn=post_softmax_attn,
|
|
)
|
|
|
|
return self.to_out(out), intermediates
|
|
|
|
|
|
class AttentionLayers(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
depth,
|
|
heads=8,
|
|
causal=False,
|
|
cross_attend=False,
|
|
only_cross=False,
|
|
use_scalenorm=False,
|
|
use_rmsnorm=False,
|
|
use_rezero=False,
|
|
rel_pos_num_buckets=32,
|
|
rel_pos_max_distance=128,
|
|
position_infused_attn=False,
|
|
custom_layers=None,
|
|
sandwich_coef=None,
|
|
par_ratio=None,
|
|
residual_attn=False,
|
|
cross_residual_attn=False,
|
|
macaron=False,
|
|
pre_norm=True,
|
|
gate_residual=False,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
ff_kwargs, kwargs = groupby_prefix_and_trim("ff_", kwargs)
|
|
attn_kwargs, _ = groupby_prefix_and_trim("attn_", kwargs)
|
|
|
|
dim_head = attn_kwargs.get("dim_head", DEFAULT_DIM_HEAD)
|
|
|
|
self.dim = dim
|
|
self.depth = depth
|
|
self.layers = nn.ModuleList([])
|
|
|
|
self.has_pos_emb = position_infused_attn
|
|
self.pia_pos_emb = (
|
|
FixedPositionalEmbedding(dim) if position_infused_attn else None
|
|
)
|
|
self.rotary_pos_emb = always(None)
|
|
|
|
assert (
|
|
rel_pos_num_buckets <= rel_pos_max_distance
|
|
), "number of relative position buckets must be less than the relative position max distance"
|
|
self.rel_pos = None
|
|
|
|
self.pre_norm = pre_norm
|
|
|
|
self.residual_attn = residual_attn
|
|
self.cross_residual_attn = cross_residual_attn
|
|
|
|
norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
|
|
norm_class = RMSNorm if use_rmsnorm else norm_class
|
|
norm_fn = partial(norm_class, dim)
|
|
|
|
norm_fn = nn.Identity if use_rezero else norm_fn
|
|
branch_fn = Rezero if use_rezero else None
|
|
|
|
if cross_attend and not only_cross:
|
|
default_block = ("a", "c", "f")
|
|
elif cross_attend and only_cross:
|
|
default_block = ("c", "f")
|
|
else:
|
|
default_block = ("a", "f")
|
|
|
|
if macaron:
|
|
default_block = ("f",) + default_block
|
|
|
|
if exists(custom_layers):
|
|
layer_types = custom_layers
|
|
elif exists(par_ratio):
|
|
par_depth = depth * len(default_block)
|
|
assert 1 < par_ratio <= par_depth, "par ratio out of range"
|
|
default_block = tuple(filter(not_equals("f"), default_block))
|
|
par_attn = par_depth // par_ratio
|
|
depth_cut = (
|
|
par_depth * 2 // 3
|
|
) # 2 / 3 attention layer cutoff suggested by PAR paper
|
|
par_width = (depth_cut + depth_cut // par_attn) // par_attn
|
|
assert (
|
|
len(default_block) <= par_width
|
|
), "default block is too large for par_ratio"
|
|
par_block = default_block + ("f",) * (par_width - len(default_block))
|
|
par_head = par_block * par_attn
|
|
layer_types = par_head + ("f",) * (par_depth - len(par_head))
|
|
elif exists(sandwich_coef):
|
|
assert (
|
|
sandwich_coef > 0 and sandwich_coef <= depth
|
|
), "sandwich coefficient should be less than the depth"
|
|
layer_types = (
|
|
("a",) * sandwich_coef
|
|
+ default_block * (depth - sandwich_coef)
|
|
+ ("f",) * sandwich_coef
|
|
)
|
|
else:
|
|
layer_types = default_block * depth
|
|
|
|
self.layer_types = layer_types
|
|
self.num_attn_layers = len(list(filter(equals("a"), layer_types)))
|
|
|
|
for layer_type in self.layer_types:
|
|
if layer_type == "a":
|
|
layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
|
|
elif layer_type == "c":
|
|
layer = Attention(dim, heads=heads, **attn_kwargs)
|
|
elif layer_type == "f":
|
|
layer = FeedForward(dim, **ff_kwargs)
|
|
layer = layer if not macaron else Scale(0.5, layer)
|
|
else:
|
|
raise Exception(f"invalid layer type {layer_type}")
|
|
|
|
if isinstance(layer, Attention) and exists(branch_fn):
|
|
layer = branch_fn(layer)
|
|
|
|
if gate_residual:
|
|
residual_fn = GRUGating(dim)
|
|
else:
|
|
residual_fn = Residual()
|
|
|
|
self.layers.append(nn.ModuleList([norm_fn(), layer, residual_fn]))
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
context=None,
|
|
mask=None,
|
|
context_mask=None,
|
|
mems=None,
|
|
return_hiddens=False,
|
|
**kwargs,
|
|
):
|
|
hiddens = []
|
|
intermediates = []
|
|
prev_attn = None
|
|
prev_cross_attn = None
|
|
|
|
mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
|
|
|
|
for ind, (layer_type, (norm, block, residual_fn)) in enumerate(
|
|
zip(self.layer_types, self.layers)
|
|
):
|
|
is_last = ind == (len(self.layers) - 1)
|
|
|
|
if layer_type == "a":
|
|
hiddens.append(x)
|
|
layer_mem = mems.pop(0)
|
|
|
|
residual = x
|
|
|
|
if self.pre_norm:
|
|
x = norm(x)
|
|
|
|
if layer_type == "a":
|
|
out, inter = block(
|
|
x,
|
|
mask=mask,
|
|
sinusoidal_emb=self.pia_pos_emb,
|
|
rel_pos=self.rel_pos,
|
|
prev_attn=prev_attn,
|
|
mem=layer_mem,
|
|
)
|
|
elif layer_type == "c":
|
|
out, inter = block(
|
|
x,
|
|
context=context,
|
|
mask=mask,
|
|
context_mask=context_mask,
|
|
prev_attn=prev_cross_attn,
|
|
)
|
|
elif layer_type == "f":
|
|
out = block(x)
|
|
|
|
x = residual_fn(out, residual)
|
|
|
|
if layer_type in ("a", "c"):
|
|
intermediates.append(inter)
|
|
|
|
if layer_type == "a" and self.residual_attn:
|
|
prev_attn = inter.pre_softmax_attn
|
|
elif layer_type == "c" and self.cross_residual_attn:
|
|
prev_cross_attn = inter.pre_softmax_attn
|
|
|
|
if not self.pre_norm and not is_last:
|
|
x = norm(x)
|
|
|
|
if return_hiddens:
|
|
intermediates = LayerIntermediates(
|
|
hiddens=hiddens, attn_intermediates=intermediates
|
|
)
|
|
|
|
return x, intermediates
|
|
|
|
return x
|
|
|
|
|
|
class Encoder(AttentionLayers):
|
|
def __init__(self, **kwargs):
|
|
assert "causal" not in kwargs, "cannot set causality on encoder"
|
|
super().__init__(causal=False, **kwargs)
|
|
|
|
|
|
class TransformerWrapper(nn.Module):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
num_tokens,
|
|
max_seq_len,
|
|
attn_layers,
|
|
emb_dim=None,
|
|
max_mem_len=0.0,
|
|
emb_dropout=0.0,
|
|
num_memory_tokens=None,
|
|
tie_embedding=False,
|
|
use_pos_emb=True,
|
|
):
|
|
super().__init__()
|
|
assert isinstance(
|
|
attn_layers, AttentionLayers
|
|
), "attention layers must be one of Encoder or Decoder"
|
|
|
|
dim = attn_layers.dim
|
|
emb_dim = default(emb_dim, dim)
|
|
|
|
self.max_seq_len = max_seq_len
|
|
self.max_mem_len = max_mem_len
|
|
self.num_tokens = num_tokens
|
|
|
|
self.token_emb = nn.Embedding(num_tokens, emb_dim)
|
|
self.pos_emb = (
|
|
AbsolutePositionalEmbedding(emb_dim, max_seq_len)
|
|
if (use_pos_emb and not attn_layers.has_pos_emb)
|
|
else always(0)
|
|
)
|
|
self.emb_dropout = nn.Dropout(emb_dropout)
|
|
|
|
self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
|
|
self.attn_layers = attn_layers
|
|
self.norm = nn.LayerNorm(dim)
|
|
|
|
self.init_()
|
|
|
|
self.to_logits = (
|
|
nn.Linear(dim, num_tokens)
|
|
if not tie_embedding
|
|
else lambda t: t @ self.token_emb.weight.t()
|
|
)
|
|
|
|
# memory tokens (like [cls]) from Memory Transformers paper
|
|
num_memory_tokens = default(num_memory_tokens, 0)
|
|
self.num_memory_tokens = num_memory_tokens
|
|
if num_memory_tokens > 0:
|
|
self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
|
|
|
|
# let funnel encoder know number of memory tokens, if specified
|
|
if hasattr(attn_layers, "num_memory_tokens"):
|
|
attn_layers.num_memory_tokens = num_memory_tokens
|
|
|
|
def init_(self):
|
|
nn.init.normal_(self.token_emb.weight, std=0.02)
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
return_embeddings=False,
|
|
mask=None,
|
|
return_mems=False,
|
|
return_attn=False,
|
|
mems=None,
|
|
embedding_manager=None,
|
|
**kwargs,
|
|
):
|
|
b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens
|
|
|
|
embedded_x = self.token_emb(x)
|
|
|
|
if embedding_manager:
|
|
x = embedding_manager(x, embedded_x)
|
|
else:
|
|
x = embedded_x
|
|
|
|
x = x + self.pos_emb(x)
|
|
x = self.emb_dropout(x)
|
|
|
|
x = self.project_emb(x)
|
|
|
|
if num_mem > 0:
|
|
mem = repeat(self.memory_tokens, "n d -> b n d", b=b)
|
|
x = torch.cat((mem, x), dim=1)
|
|
|
|
# auto-handle masking after appending memory tokens
|
|
if exists(mask):
|
|
mask = F.pad(mask, (num_mem, 0), value=True)
|
|
|
|
x, intermediates = self.attn_layers(
|
|
x, mask=mask, mems=mems, return_hiddens=True, **kwargs
|
|
)
|
|
x = self.norm(x)
|
|
|
|
mem, x = x[:, :num_mem], x[:, num_mem:]
|
|
|
|
out = self.to_logits(x) if not return_embeddings else x
|
|
|
|
if return_mems:
|
|
hiddens = intermediates.hiddens
|
|
new_mems = (
|
|
list(
|
|
map(
|
|
lambda pair: torch.cat(pair, dim=-2),
|
|
zip(mems, hiddens),
|
|
)
|
|
)
|
|
if exists(mems)
|
|
else hiddens
|
|
)
|
|
new_mems = list(
|
|
map(lambda t: t[..., -self.max_mem_len :, :].detach(), new_mems)
|
|
)
|
|
return out, new_mems
|
|
|
|
if return_attn:
|
|
attn_maps = list(
|
|
map(
|
|
lambda t: t.post_softmax_attn,
|
|
intermediates.attn_intermediates,
|
|
)
|
|
)
|
|
return out, attn_maps
|
|
|
|
return out
|