File size: 2,456 Bytes
484364b | 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 | import torch, torch.nn as nn, math
class M(nn.Module):
def __init__(self, d=128, L=3, H=2, nc=10):
super().__init__()
self.pe = nn.Conv2d(3, d, 16, 16)
self.cls = nn.Parameter(torch.randn(1,1,d)*.02)
self.pos = nn.Parameter(torch.randn(1,197,d)*.02)
self.blks = nn.ModuleList([nn.TransformerEncoderLayer(d, H, d*4, .1, activation='gelu', batch_first=True, norm_first=True) for _ in range(L)])
self.ln = nn.LayerNorm(d)
self.te = nn.Embedding(30522, d, padding_idx=0)
self.tpos = nn.Parameter(torch.randn(1,128,d)*.02)
self.tblks = nn.ModuleList([nn.TransformerEncoderLayer(d, H, d*4, .1, activation='gelu', batch_first=True, norm_first=True) for _ in range(L)])
self.tln = nn.LayerNorm(d)
self.fuse = nn.ModuleList([nn.TransformerEncoderLayer(d, H, d*4, .1, activation='gelu', batch_first=True, norm_first=True) for _ in range(2)])
self.fln = nn.LayerNorm(d)
self.head = nn.Sequential(nn.Linear(d,d), nn.GELU(approximate='quick'), nn.Dropout(.1), nn.Linear(d,nc))
self._init()
def _init(self):
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.trunc_normal_(m.weight, std=0.02)
if m.bias is not None: nn.init.zeros_(m.bias)
def enc_img(self, x):
x = self.pe(x).flatten(2).transpose(1,2)
x = torch.cat([self.cls.expand(x.size(0),-1,-1), x], 1)
x = x + self.pos
for b in self.blks: x = b(x)
return self.ln(x)
def enc_txt(self, ids, mask=None):
x = self.te(ids) + self.tpos[:, :ids.size(1)]
m = (mask == 0) if mask is not None else None
for b in self.tblks: x = b(x, src_key_padding_mask=m)
return self.tln(x)
def forward(self, img, ids, mask=None, lbl=None):
fi = self.enc_img(img)
ft = self.enc_txt(ids, mask)
x = ft
for f in self.fuse: x = f(x)
x = self.fln(x[:, 0])
logits = self.head(x)
loss = nn.functional.cross_entropy(logits, lbl) if lbl is not None else None
return {'logits': logits, 'loss': loss}
if __name__ == '__main__':
m = M()
print(f'Params: {sum(p.numel() for p in m.parameters()):,}')
o = m(torch.randn(2,3,224,224), torch.randint(0,30522,(2,128)), torch.ones(2,128), torch.tensor([0,1]))
print(o['logits'].shape, o['loss'].item())
|