Spaces:
Runtime error
Runtime error
File size: 3,343 Bytes
3672502 |
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 |
"""
Copied from
https://github.com/lucidrains/flamingo-pytorch/blob/main/flamingo_pytorch/flamingo_pytorch.py
"""
import torch
from torch import nn, einsum
import torch.nn.functional as F
from einops import rearrange, repeat
from einops_exts import rearrange_many, repeat_many
def exists(val):
return val is not None
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,
*,
dim,
dim_head = 64,
heads = 8
):
super().__init__()
self.scale = dim_head ** -0.5
self.heads = heads
inner_dim = dim_head * heads
self.norm_media = nn.LayerNorm(dim)
self.norm_latents = nn.LayerNorm(dim)
self.to_q = nn.Linear(dim, inner_dim, bias = False)
self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
self.to_out = nn.Linear(inner_dim, dim, 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)
# the paper differs from Perceiver in which they also concat the key / values derived from the latents to be attended to
kv_input = torch.cat((x, latents), dim = -2)
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,
*,
dim,
depth,
dim_head = 64,
heads = 8,
num_latents = 64,
num_media_embeds = 4,
ff_mult = 4
):
super().__init__()
self.latents = nn.Parameter(torch.randn(num_latents, dim))
self.media_pos_emb = nn.Parameter(torch.randn(num_media_embeds, 1, dim))
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
PerceiverAttention(dim = dim, dim_head = dim_head, heads = heads),
FeedForward(dim = dim, mult = ff_mult)
]))
self.norm = nn.LayerNorm(dim)
def forward(self, x):
if x.ndim == 3:
x = rearrange(x, 'b n d -> b 1 n d')
times = x.shape[1]
x = x + self.media_pos_emb[:times]
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
res = self.norm(latents)
if res.ndim == 4:
res = res.squeeze(1)
return res |