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