| 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) |
|
|
