File size: 2,970 Bytes
cb5d220 | 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 | import torch
import torch.nn as nn
import torch.nn.functional as F
D = 128
D2 = 96
class DepthwiseSeparableConv(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.dw = nn.Conv2d(in_ch, in_ch, 3, padding=1, groups=in_ch)
self.pw = nn.Conv2d(in_ch, out_ch, 1)
def forward(self, x):
return self.pw(self.dw(x))
class AtomEncoder(nn.Module):
def __init__(self, embed_dim=D):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.GELU(),
nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.GELU(),
DepthwiseSeparableConv(32, 64), nn.GELU(),
nn.AdaptiveAvgPool2d(1),
)
self.proj = nn.Linear(64, embed_dim)
def forward(self, x):
return self.proj(self.net(x).flatten(1))
class GridRelationalAttention(nn.Module):
def __init__(self, dim=D, heads=4, n_relations=4):
super().__init__()
self.missing_token = nn.Parameter(torch.randn(dim) * 0.02)
self.row_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
self.col_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
self.relation_queries = nn.Parameter(torch.randn(n_relations, dim) * 0.02)
self.cross_attn = nn.MultiheadAttention(dim, heads, batch_first=True)
self.rule_proj = nn.Linear(n_relations * dim, dim)
self.n_relations, self.dim = n_relations, dim
def forward(self, context_panels):
B = context_panels.shape[0]
missing = self.missing_token.expand(B, 1, self.dim)
grid = torch.cat([context_panels, missing], dim=1)
rows_idx = [[0,1,2],[3,4,5],[6,7,8]]
cols_idx = [[0,3,6],[1,4,7],[2,5,8]]
row_vecs = torch.cat([
self.row_attn(grid[:, idx, :], grid[:, idx, :], grid[:, idx, :])[0].mean(1, keepdim=True)
for idx in rows_idx
], dim=1)
col_vecs = torch.cat([
self.col_attn(grid[:, idx, :], grid[:, idx, :], grid[:, idx, :])[0].mean(1, keepdim=True)
for idx in cols_idx
], dim=1)
combined_ctx = torch.cat([row_vecs, col_vecs], dim=1)
rq = self.relation_queries.unsqueeze(0).expand(B, -1, -1)
rel_out, _ = self.cross_attn(rq, combined_ctx, combined_ctx)
return self.rule_proj(rel_out.flatten(1))
class CandidateScorer(nn.Module):
def __init__(self, dim=D):
super().__init__()
self.score_proj = nn.Linear(dim, dim)
def forward(self, rule_embedding, candidate_embeds):
proj_rule = self.score_proj(rule_embedding)
return torch.einsum("bd,bkd->bk", proj_rule, candidate_embeds)
class GRAFT(nn.Module):
def __init__(self, vocab_size, narrator_cls):
super().__init__()
self.encoder = AtomEncoder()
self.gra = GridRelationalAttention()
self.scorer = CandidateScorer()
self.narrator = narrator_cls(vocab_size=vocab_size)
|