Spaces:
Running
Running
File size: 1,941 Bytes
ac59957 |
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 |
# https://github.com/XiaoyuShi97/VideoFlow/blob/main/core/Networks/BOFNet/gma.py
import torch
import math
from torch import nn, einsum
from einops import rearrange
class Attention(nn.Module):
def __init__(
self,
dim: int,
heads: int,
dim_head: int,
):
super().__init__()
self.heads = heads
self.scale = dim_head**-0.5
inner_dim = heads * dim_head
self.to_qk = nn.Conv2d(dim, inner_dim * 2, 1, bias=False)
def forward(self, fmap):
heads, _, _, h, w = self.heads, *fmap.shape
q, k = self.to_qk(fmap).chunk(2, dim=1)
q, k = map(lambda t: rearrange(t, "b (h d) x y -> b h x y d", h=heads), (q, k))
# Small change based on MemFlow Paper
q = self.scale * q * math.log(h * w, 3)
sim = einsum("b h x y d, b h u v d -> b h x y u v", q, k)
sim = rearrange(sim, "b h x y u v -> b h (x y) (u v)")
attn = sim.softmax(dim=-1)
return attn
class Aggregate(nn.Module):
def __init__(
self,
dim,
heads=4,
dim_head=128,
):
super().__init__()
self.heads = heads
self.scale = dim_head**-0.5
inner_dim = heads * dim_head
self.to_v = nn.Conv2d(dim, inner_dim, 1, bias=False)
self.gamma = nn.Parameter(torch.zeros(1))
if dim != inner_dim:
self.project = nn.Conv2d(inner_dim, dim, 1, bias=False)
else:
self.project = None
def forward(self, attn, fmap):
heads, _, _, h, w = self.heads, *fmap.shape
v = self.to_v(fmap)
v = rearrange(v, "b (h d) x y -> b h (x y) d", h=heads)
out = einsum("b h i j, b h j d -> b h i d", attn, v)
out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
if self.project is not None:
out = self.project(out)
out = fmap + self.gamma * out
return out
|