|
| 1 | +""" |
| 2 | +https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit_pytorch.py |
| 3 | +""" |
| 4 | + |
| 5 | +import torch |
| 6 | +import torch.nn.functional as F |
| 7 | +from einops import rearrange, repeat |
| 8 | +from torch import nn |
| 9 | + |
| 10 | +MIN_NUM_PATCHES = 16 |
| 11 | + |
| 12 | +class Residual(nn.Module): |
| 13 | + def __init__(self, fn): |
| 14 | + super().__init__() |
| 15 | + self.fn = fn |
| 16 | + def forward(self, x, **kwargs): |
| 17 | + return self.fn(x, **kwargs) + x |
| 18 | + |
| 19 | +class PreNorm(nn.Module): |
| 20 | + def __init__(self, dim, fn): |
| 21 | + super().__init__() |
| 22 | + self.norm = nn.LayerNorm(dim) |
| 23 | + self.fn = fn |
| 24 | + def forward(self, x, **kwargs): |
| 25 | + return self.fn(self.norm(x), **kwargs) |
| 26 | + |
| 27 | +class FeedForward(nn.Module): |
| 28 | + def __init__(self, dim, hidden_dim, dropout = 0.): |
| 29 | + super().__init__() |
| 30 | + self.net = nn.Sequential( |
| 31 | + nn.Linear(dim, hidden_dim), |
| 32 | + nn.GELU(), |
| 33 | + nn.Dropout(dropout), |
| 34 | + nn.Linear(hidden_dim, dim), |
| 35 | + nn.Dropout(dropout) |
| 36 | + ) |
| 37 | + def forward(self, x): |
| 38 | + return self.net(x) |
| 39 | + |
| 40 | +class Attention(nn.Module): |
| 41 | + def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.): |
| 42 | + super().__init__() |
| 43 | + inner_dim = dim_head * heads |
| 44 | + self.heads = heads |
| 45 | + self.scale = dim_head ** -0.5 |
| 46 | + |
| 47 | + self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) |
| 48 | + self.to_out = nn.Sequential( |
| 49 | + nn.Linear(inner_dim, dim), |
| 50 | + nn.Dropout(dropout) |
| 51 | + ) |
| 52 | + |
| 53 | + def forward(self, x, mask = None): |
| 54 | + b, n, _, h = *x.shape, self.heads |
| 55 | + qkv = self.to_qkv(x).chunk(3, dim = -1) |
| 56 | + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv) |
| 57 | + |
| 58 | + dots = torch.einsum('bhid,bhjd->bhij', q, k) * self.scale |
| 59 | + mask_value = -torch.finfo(dots.dtype).max |
| 60 | + |
| 61 | + if mask is not None: |
| 62 | + mask = F.pad(mask.flatten(1), (1, 0), value = True) |
| 63 | + assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions' |
| 64 | + mask = mask[:, None, :] * mask[:, :, None] |
| 65 | + dots.masked_fill_(~mask, mask_value) |
| 66 | + del mask |
| 67 | + |
| 68 | + attn = dots.softmax(dim=-1) |
| 69 | + |
| 70 | + out = torch.einsum('bhij,bhjd->bhid', attn, v) |
| 71 | + out = rearrange(out, 'b h n d -> b n (h d)') |
| 72 | + out = self.to_out(out) |
| 73 | + return out |
| 74 | + |
| 75 | +class Transformer(nn.Module): |
| 76 | + def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout): |
| 77 | + super().__init__() |
| 78 | + self.layers = nn.ModuleList([]) |
| 79 | + for _ in range(depth): |
| 80 | + self.layers.append(nn.ModuleList([ |
| 81 | + Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))), |
| 82 | + Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))) |
| 83 | + ])) |
| 84 | + def forward(self, x, mask = None): |
| 85 | + for attn, ff in self.layers: |
| 86 | + x = attn(x, mask = mask) |
| 87 | + x = ff(x) |
| 88 | + return x |
| 89 | + |
| 90 | +class ViT(nn.Module): |
| 91 | + def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.): |
| 92 | + super().__init__() |
| 93 | + assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.' |
| 94 | + num_patches = (image_size // patch_size) ** 2 |
| 95 | + patch_dim = channels * patch_size ** 2 |
| 96 | + assert num_patches > MIN_NUM_PATCHES, f'your number of patches ({num_patches}) is way too small for attention to be effective (at least 16). Try decreasing your patch size' |
| 97 | + assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)' |
| 98 | + |
| 99 | + self.patch_size = patch_size |
| 100 | + |
| 101 | + self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim)) |
| 102 | + self.patch_to_embedding = nn.Linear(patch_dim, dim) |
| 103 | + self.cls_token = nn.Parameter(torch.randn(1, 1, dim)) |
| 104 | + self.dropout = nn.Dropout(emb_dropout) |
| 105 | + |
| 106 | + self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout) |
| 107 | + |
| 108 | + self.pool = pool |
| 109 | + self.to_latent = nn.Identity() |
| 110 | + |
| 111 | + self.mlp_head = nn.Sequential( |
| 112 | + nn.LayerNorm(dim), |
| 113 | + nn.Linear(dim, num_classes) |
| 114 | + ) |
| 115 | + |
| 116 | + def forward(self, img, mask = None): |
| 117 | + p = self.patch_size |
| 118 | + |
| 119 | + x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p) |
| 120 | + x = self.patch_to_embedding(x) |
| 121 | + b, n, _ = x.shape |
| 122 | + |
| 123 | + cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b) |
| 124 | + x = torch.cat((cls_tokens, x), dim=1) |
| 125 | + x += self.pos_embedding[:, :(n + 1)] |
| 126 | + x = self.dropout(x) |
| 127 | + |
| 128 | + x = self.transformer(x, mask) |
| 129 | + |
| 130 | + x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0] |
| 131 | + |
| 132 | + x = self.to_latent(x) |
| 133 | + return self.mlp_head(x) |
0 commit comments