|
| 1 | +import torch |
| 2 | +from torch import nn |
| 3 | + |
| 4 | +from einops import rearrange, repeat |
| 5 | +from einops.layers.torch import Rearrange |
| 6 | + |
| 7 | +# helpers |
| 8 | + |
| 9 | +def pair(t): |
| 10 | + return t if isinstance(t, tuple) else (t, t) |
| 11 | + |
| 12 | +# classes |
| 13 | + |
| 14 | +class Parallel(nn.Module): |
| 15 | + def __init__(self, *fns): |
| 16 | + super().__init__() |
| 17 | + self.fns = nn.ModuleList(fns) |
| 18 | + |
| 19 | + def forward(self, x): |
| 20 | + return sum([fn(x) for fn in self.fns]) |
| 21 | + |
| 22 | +class PreNorm(nn.Module): |
| 23 | + def __init__(self, dim, fn): |
| 24 | + super().__init__() |
| 25 | + self.norm = nn.LayerNorm(dim) |
| 26 | + self.fn = fn |
| 27 | + def forward(self, x, **kwargs): |
| 28 | + return self.fn(self.norm(x), **kwargs) |
| 29 | + |
| 30 | +class FeedForward(nn.Module): |
| 31 | + def __init__(self, dim, hidden_dim, dropout = 0.): |
| 32 | + super().__init__() |
| 33 | + self.net = nn.Sequential( |
| 34 | + nn.Linear(dim, hidden_dim), |
| 35 | + nn.GELU(), |
| 36 | + nn.Dropout(dropout), |
| 37 | + nn.Linear(hidden_dim, dim), |
| 38 | + nn.Dropout(dropout) |
| 39 | + ) |
| 40 | + def forward(self, x): |
| 41 | + return self.net(x) |
| 42 | + |
| 43 | +class Attention(nn.Module): |
| 44 | + def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.): |
| 45 | + super().__init__() |
| 46 | + inner_dim = dim_head * heads |
| 47 | + project_out = not (heads == 1 and dim_head == dim) |
| 48 | + |
| 49 | + self.heads = heads |
| 50 | + self.scale = dim_head ** -0.5 |
| 51 | + |
| 52 | + self.attend = nn.Softmax(dim = -1) |
| 53 | + self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False) |
| 54 | + |
| 55 | + self.to_out = nn.Sequential( |
| 56 | + nn.Linear(inner_dim, dim), |
| 57 | + nn.Dropout(dropout) |
| 58 | + ) if project_out else nn.Identity() |
| 59 | + |
| 60 | + def forward(self, x): |
| 61 | + qkv = self.to_qkv(x).chunk(3, dim = -1) |
| 62 | + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv) |
| 63 | + |
| 64 | + dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale |
| 65 | + |
| 66 | + attn = self.attend(dots) |
| 67 | + |
| 68 | + out = torch.matmul(attn, v) |
| 69 | + out = rearrange(out, 'b h n d -> b n (h d)') |
| 70 | + return self.to_out(out) |
| 71 | + |
| 72 | +class Transformer(nn.Module): |
| 73 | + def __init__(self, dim, depth, heads, dim_head, mlp_dim, num_parallel_branches = 2, dropout = 0.): |
| 74 | + super().__init__() |
| 75 | + self.layers = nn.ModuleList([]) |
| 76 | + |
| 77 | + attn_block = lambda: PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)) |
| 78 | + ff_block = lambda: PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)) |
| 79 | + |
| 80 | + for _ in range(depth): |
| 81 | + self.layers.append(nn.ModuleList([ |
| 82 | + Parallel(*[attn_block() for _ in range(num_parallel_branches)]), |
| 83 | + Parallel(*[ff_block() for _ in range(num_parallel_branches)]), |
| 84 | + ])) |
| 85 | + |
| 86 | + def forward(self, x): |
| 87 | + for attns, ffs in self.layers: |
| 88 | + x = attns(x) + x |
| 89 | + x = ffs(x) + x |
| 90 | + return x |
| 91 | + |
| 92 | +class ViT(nn.Module): |
| 93 | + def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', num_parallel_branches = 2, channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.): |
| 94 | + super().__init__() |
| 95 | + image_height, image_width = pair(image_size) |
| 96 | + patch_height, patch_width = pair(patch_size) |
| 97 | + |
| 98 | + assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.' |
| 99 | + |
| 100 | + num_patches = (image_height // patch_height) * (image_width // patch_width) |
| 101 | + patch_dim = channels * patch_height * patch_width |
| 102 | + assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)' |
| 103 | + |
| 104 | + self.to_patch_embedding = nn.Sequential( |
| 105 | + Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width), |
| 106 | + nn.Linear(patch_dim, dim), |
| 107 | + ) |
| 108 | + |
| 109 | + self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim)) |
| 110 | + self.cls_token = nn.Parameter(torch.randn(1, 1, dim)) |
| 111 | + self.dropout = nn.Dropout(emb_dropout) |
| 112 | + |
| 113 | + self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, num_parallel_branches, dropout) |
| 114 | + |
| 115 | + self.pool = pool |
| 116 | + self.to_latent = nn.Identity() |
| 117 | + |
| 118 | + self.mlp_head = nn.Sequential( |
| 119 | + nn.LayerNorm(dim), |
| 120 | + nn.Linear(dim, num_classes) |
| 121 | + ) |
| 122 | + |
| 123 | + def forward(self, img): |
| 124 | + x = self.to_patch_embedding(img) |
| 125 | + b, n, _ = x.shape |
| 126 | + |
| 127 | + cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b) |
| 128 | + x = torch.cat((cls_tokens, x), dim=1) |
| 129 | + x += self.pos_embedding[:, :(n + 1)] |
| 130 | + x = self.dropout(x) |
| 131 | + |
| 132 | + x = self.transformer(x) |
| 133 | + |
| 134 | + x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0] |
| 135 | + |
| 136 | + x = self.to_latent(x) |
| 137 | + return self.mlp_head(x) |
0 commit comments