File size: 3,392 Bytes
cee9fbc |
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 |
# This file may have been modified by Bytedance Ltd. and/or its affiliates (“Bytedance's Modifications”).
# All Bytedance's Modifications are Copyright (year) Bytedance Ltd. and/or its affiliates.
import torch
from torch import nn, einsum
from einops import rearrange, repeat
from einops_exts import rearrange_many, repeat_many
def FeedForward(dim, mult=4):
inner_dim = int(dim * mult)
return nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, inner_dim, bias=False),
nn.GELU(),
nn.Linear(inner_dim, dim, bias=False)
)
class PerceiverAttention(nn.Module):
def __init__(
self,
vision_width,
text_width,
dim_head=64,
heads=8
):
super().__init__()
self.vision_width = vision_width
self.text_width = text_width
self.scale = dim_head ** -0.5
self.heads = heads
inner_dim = dim_head * heads
self.norm_media = nn.LayerNorm(vision_width)
self.norm_latents = nn.LayerNorm(text_width)
self.to_q = nn.Linear(text_width, inner_dim, bias=False)
self.to_kv = nn.Linear(vision_width, inner_dim * 2, bias=False)
self.to_out = nn.Linear(inner_dim, text_width, bias=False)
def forward(self, x, latents):
"""
einstein notation
b - batch
t - time
n - sequence
d - dimension
"""
x = self.norm_media(x)
latents = self.norm_latents(latents)
b, m, h = *x.shape[:2], self.heads
q = self.to_q(latents)
kv_input = x
k, v = self.to_kv(kv_input).chunk(2, dim=-1)
q, k, v = rearrange_many((q, k, v), 'b t n (h d) -> b h t n d', h=h)
q = q * self.scale
# attention
sim = einsum('... i d, ... j d -> ... i j', q, k)
sim = sim - sim.amax(dim=-1, keepdim=True).detach()
attn = sim.softmax(dim=-1)
out = einsum('... i j, ... j d -> ... i d', attn, v)
out = rearrange(out, 'b h t n d -> b t n (h d)', h=h)
return self.to_out(out)
class PerceiverResampler(nn.Module):
def __init__(
self,
vision_width,
text_width,
depth,
dim_head=64,
heads=8,
num_latents=64,
ff_mult=4,
):
super().__init__()
self.latents = nn.Parameter(torch.randn(num_latents, text_width))
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
PerceiverAttention(vision_width=vision_width, text_width=text_width, dim_head=dim_head, heads=heads),
FeedForward(dim=text_width, mult=ff_mult)
]))
self.norm = nn.LayerNorm(text_width)
def forward(self, vision_embeds=None, vision_atts=None):
x = vision_embeds
if x.ndim == 3:
x = rearrange(x, 'b n d -> b 1 n d')
latents = repeat(self.latents, 'n d -> b m n d', b=x.shape[0], m=x.shape[1])
for attn, ff in self.layers:
latents = attn(x, latents) + latents
latents = ff(latents) + latents
v2t_feats = self.norm(latents).squeeze(dim=1) # for image, squeeze dim=1
v2t_atts = torch.ones(v2t_feats.shape[:2], dtype=torch.long, device=v2t_feats.device)
return v2t_feats, v2t_atts |