add CrossViT

Masked Patch Prediction "Suggested in #63" Work in Progress

README.md

@@ -1,4 +1,4 @@
-<img src="./vit.gif" width="500px"></img>
+<img src="./images/vit.gif" width="500px"></img>
 
 ## Vision Transformer - Pytorch
 
@@ -33,9 +33,8 @@ v = ViT(
 )
 
 img = torch.randn(1, 3, 256, 256)
-mask = torch.ones(1, 8, 8).bool() # optional mask, designating which patch to attend to
 
-preds = v(img, mask = mask) # (1, 1000)
+preds = v(img) # (1, 1000)
 ```
 
 ## Parameters
@@ -64,7 +63,7 @@ Embedding dropout rate.
 
 ## Distillation
 
-<img src="./distill.png" width="300px"></img>
+<img src="./images/distill.png" width="300px"></img>
 
 A recent <a href="https://arxiv.org/abs/2012.12877">paper</a> has shown that use of a distillation token for distilling knowledge from convolutional nets to vision transformer can yield small and efficient vision transformers. This repository offers the means to do distillation easily.
 
@@ -117,9 +116,36 @@ v = v.to_vit()
 type(v) # <class 'vit_pytorch.vit_pytorch.ViT'>
 ```
 
+## Deep ViT
+
+This <a href="https://arxiv.org/abs/2103.11886">paper</a> notes that ViT struggles to attend at greater depths (past 12 layers), and suggests mixing the attention of each head post-softmax as a solution, dubbed Re-attention. The results line up with the <a href="https://github.com/lucidrains/x-transformers#talking-heads-attention">Talking Heads</a> paper from NLP.
+
+You can use it as follows
+
+```python
+import torch
+from vit_pytorch.deepvit import DeepViT
+
+v = DeepViT(
+    image_size = 256,
+    patch_size = 32,
+    num_classes = 1000,
+    dim = 1024,
+    depth = 6,
+    heads = 16,
+    mlp_dim = 2048,
+    dropout = 0.1,
+    emb_dropout = 0.1
+)
+
+img = torch.randn(1, 3, 256, 256)
+
+preds = v(img) # (1, 1000)
+```
+
 ## Token-to-Token ViT
 
-<img src="./t2t.png" width="400px"></img>
+<img src="./images/t2t.png" width="400px"></img>
 
 <a href="https://arxiv.org/abs/2101.11986">This paper</a> proposes that the first couple layers should downsample the image sequence by unfolding, leading to overlapping image data in each token as shown in the figure above. You can use this variant of the `ViT` as follows.
 
@@ -141,6 +167,93 @@ img = torch.randn(1, 3, 224, 224)
 v(img) # (1, 1000)
 ```
 
+## Masked Patch Prediction
+
+Thanks to <a href="https://github.com/zankner">Zach</a>, you can train using the original masked patch prediction task presented in the paper, with the following code.
+
+```python
+import torch
+from vit_pytorch import ViT
+from vit_pytorch.mpp import MPP
+
+model = ViT(
+    image_size=256,
+    patch_size=32,
+    num_classes=1000,
+    dim=1024,
+    depth=6,
+    heads=8,
+    mlp_dim=2048,
+    dropout=0.1,
+    emb_dropout=0.1
+)
+
+mpp_trainer = MPP(
+    transformer=model,
+    patch_size=32,
+    dim=1024,
+    mask_prob=0.15,          # probability of using token in masked prediction task
+    random_patch_prob=0.30,  # probability of randomly replacing a token being used for mpp
+    replace_prob=0.50,       # probability of replacing a token being used for mpp with the mask token
+)
+
+opt = torch.optim.Adam(mpp_trainer.parameters(), lr=3e-4)
+
+def sample_unlabelled_images():
+    return torch.randn(20, 3, 256, 256)
+
+for _ in range(100):
+    images = sample_unlabelled_images()
+    loss = mpp_trainer(images)
+    opt.zero_grad()
+    loss.backward()
+    opt.step()
+
+# save your improved network
+torch.save(model.state_dict(), './pretrained-net.pt')
+```
+
+## Accessing Attention
+
+If you would like to visualize the attention weights (post-softmax) for your research, just follow the procedure below
+
+```python
+import torch
+from vit_pytorch.vit import ViT
+
+v = ViT(
+    image_size = 256,
+    patch_size = 32,
+    num_classes = 1000,
+    dim = 1024,
+    depth = 6,
+    heads = 16,
+    mlp_dim = 2048,
+    dropout = 0.1,
+    emb_dropout = 0.1
+)
+
+# import Recorder and wrap the ViT
+
+from vit_pytorch.recorder import Recorder
+v = Recorder(v)
+
+# forward pass now returns predictions and the attention maps
+
+img = torch.randn(1, 3, 256, 256)
+preds, attns = v(img)
+
+# there is one extra patch due to the CLS token
+
+attns # (1, 6, 16, 65, 65) - (batch x layers x heads x patch x patch)
+```
+
+to cleanup the class and the hooks once you have collected enough data
+
+```python
+v = v.eject()  # wrapper is discarded and original ViT instance is returned
+```
+
 ## Research Ideas
 
 ### Self Supervised Training
@@ -299,12 +412,34 @@ Coming from computer vision and new to transformers? Here are some resources tha
 
 ```bibtex
 @misc{yuan2021tokenstotoken,
-      title     = {Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet},
-      author    = {Li Yuan and Yunpeng Chen and Tao Wang and Weihao Yu and Yujun Shi and Francis EH Tay and Jiashi Feng and Shuicheng Yan},
-      year      = {2021},
-      eprint    = {2101.11986},
-      archivePrefix = {arXiv},
-      primaryClass = {cs.CV}
+    title     = {Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet},
+    author    = {Li Yuan and Yunpeng Chen and Tao Wang and Weihao Yu and Yujun Shi and Francis EH Tay and Jiashi Feng and Shuicheng Yan},
+    year      = {2021},
+    eprint    = {2101.11986},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV}
+}
+```
+
+```bibtex
+@misc{zhou2021deepvit,
+    title   = {DeepViT: Towards Deeper Vision Transformer},
+    author  = {Daquan Zhou and Bingyi Kang and Xiaojie Jin and Linjie Yang and Xiaochen Lian and Qibin Hou and Jiashi Feng},
+    year    = {2021},
+    eprint  = {2103.11886},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV}
+}
+```
+
+```bibtex
+@misc{chen2021crossvit,
+    title   = {CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification},
+    author  = {Chun-Fu Chen and Quanfu Fan and Rameswar Panda},
+    year    = {2021},
+    eprint  = {2103.14899},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.CV}
 }
 ```
 

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'vit-pytorch',
   packages = find_packages(exclude=['examples']),
-  version = '0.7.4',
+  version = '0.10.3',
   license='MIT',
   description = 'Vision Transformer (ViT) - Pytorch',
   author = 'Phil Wang',

vit_pytorch/__init__.py

@@ -1 +1 @@
-from vit_pytorch.vit_pytorch import ViT
+from vit_pytorch.vit import ViT

vit_pytorch/cross_vit.py

@@ -0,0 +1,270 @@
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+# pre-layernorm
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+# feedforward
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+# attention
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim = -1)
+        self.to_q = nn.Linear(dim, inner_dim, bias = False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x, context = None, kv_include_self = False):
+        b, n, _, h = *x.shape, self.heads
+        context = default(context, x)
+
+        if kv_include_self:
+            context = torch.cat((x, context), dim = 1) # cross attention requires CLS token includes itself as key / value
+
+        qkv = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
+
+        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+
+        attn = self.attend(dots)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+# transformer encoder, for small and large patches
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        self.norm = nn.LayerNorm(dim)
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
+                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+            ]))
+
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+# projecting CLS tokens, in the case that small and large patch tokens have different dimensions
+
+class ProjectInOut(nn.Module):
+    def __init__(self, dim_in, dim_out, fn):
+        super().__init__()
+        self.fn = fn
+
+        need_projection = dim_in != dim_out
+        self.project_in = nn.Linear(dim_in, dim_out) if need_projection else nn.Identity()
+        self.project_out = nn.Linear(dim_out, dim_in) if need_projection else nn.Identity()
+
+    def forward(self, x, *args, **kwargs):
+        x = self.project_in(x)
+        x = self.fn(x, *args, **kwargs)
+        x = self.project_out(x)
+        return x
+
+# cross attention transformer
+
+class CrossTransformer(nn.Module):
+    def __init__(self, sm_dim, lg_dim, depth, heads, dim_head, dropout):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                ProjectInOut(sm_dim, lg_dim, PreNorm(lg_dim, Attention(lg_dim, heads = heads, dim_head = dim_head, dropout = dropout))),
+                ProjectInOut(lg_dim, sm_dim, PreNorm(sm_dim, Attention(sm_dim, heads = heads, dim_head = dim_head, dropout = dropout)))
+            ]))
+
+    def forward(self, sm_tokens, lg_tokens):
+        (sm_cls, sm_patch_tokens), (lg_cls, lg_patch_tokens) = map(lambda t: (t[:, :1], t[:, 1:]), (sm_tokens, lg_tokens))
+
+        for sm_attend_lg, lg_attend_sm in self.layers:
+            sm_cls = sm_attend_lg(sm_cls, context = lg_patch_tokens, kv_include_self = True) + sm_cls
+            lg_cls = lg_attend_sm(lg_cls, context = sm_patch_tokens, kv_include_self = True) + lg_cls
+
+        sm_tokens = torch.cat((sm_cls, sm_patch_tokens), dim = 1)
+        lg_tokens = torch.cat((lg_cls, lg_patch_tokens), dim = 1)
+        return sm_tokens, lg_tokens
+
+# multi-scale encoder
+
+class MultiScaleEncoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        depth,
+        sm_dim,
+        lg_dim,
+        sm_enc_params,
+        lg_enc_params,
+        cross_attn_heads,
+        cross_attn_depth,
+        cross_attn_dim_head = 64,
+        dropout = 0.
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Transformer(dim = sm_dim, dropout = dropout, **sm_enc_params),
+                Transformer(dim = lg_dim, dropout = dropout, **lg_enc_params),
+                CrossTransformer(sm_dim = sm_dim, lg_dim = lg_dim, depth = cross_attn_depth, heads = cross_attn_heads, dim_head = cross_attn_dim_head, dropout = dropout)
+            ]))
+
+    def forward(self, sm_tokens, lg_tokens):
+        for sm_enc, lg_enc, cross_attend in self.layers:
+            sm_tokens, lg_tokens = sm_enc(sm_tokens), lg_enc(lg_tokens)
+            sm_tokens, lg_tokens = cross_attend(sm_tokens, lg_tokens)
+
+        return sm_tokens, lg_tokens
+
+# patch-based image to token embedder
+
+class ImageEmbedder(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        image_size,
+        patch_size,
+        dropout = 0.
+    ):
+        super().__init__()
+        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
+        num_patches = (image_size // patch_size) ** 2
+        patch_dim = 3 * patch_size ** 2
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
+            nn.Linear(patch_dim, dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, img):
+        x = self.to_patch_embedding(img)
+        b, n, _ = x.shape
+
+        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding[:, :(n + 1)]
+
+        return self.dropout(x)
+
+# cross ViT class
+
+class CrossViT(nn.Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        num_classes,
+        sm_dim,
+        lg_dim,
+        sm_patch_size = 12,
+        sm_enc_depth = 1,
+        sm_enc_heads = 8,
+        sm_enc_mlp_dim = 2048,
+        sm_enc_dim_head = 64,
+        lg_patch_size = 16,
+        lg_enc_depth = 4,
+        lg_enc_heads = 8,
+        lg_enc_mlp_dim = 2048,
+        lg_enc_dim_head = 64,
+        cross_attn_depth = 2,
+        cross_attn_heads = 8,
+        cross_attn_dim_head = 64,
+        depth = 3,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    ):
+        super().__init__()
+        self.sm_image_embedder = ImageEmbedder(dim = sm_dim, image_size = image_size, patch_size = sm_patch_size, dropout = emb_dropout)
+        self.lg_image_embedder = ImageEmbedder(dim = lg_dim, image_size = image_size, patch_size = lg_patch_size, dropout = emb_dropout)
+
+        self.multi_scale_encoder = MultiScaleEncoder(
+            depth = depth,
+            sm_dim = sm_dim,
+            lg_dim = lg_dim,
+            cross_attn_heads = cross_attn_heads,
+            cross_attn_dim_head = cross_attn_dim_head,
+            cross_attn_depth = cross_attn_depth,
+            sm_enc_params = dict(
+                depth = sm_enc_depth,
+                heads = sm_enc_heads,
+                mlp_dim = sm_enc_mlp_dim,
+                dim_head = sm_enc_dim_head
+            ),
+            lg_enc_params = dict(
+                depth = lg_enc_depth,
+                heads = lg_enc_heads,
+                mlp_dim = lg_enc_mlp_dim,
+                dim_head = lg_enc_dim_head
+            ),
+            dropout = dropout
+        )
+
+        self.sm_mlp_head = nn.Sequential(nn.LayerNorm(sm_dim), nn.Linear(sm_dim, num_classes))
+        self.lg_mlp_head = nn.Sequential(nn.LayerNorm(lg_dim), nn.Linear(lg_dim, num_classes))
+
+    def forward(self, img):
+        sm_tokens = self.sm_image_embedder(img)
+        lg_tokens = self.lg_image_embedder(img)
+
+        sm_tokens, lg_tokens = self.multi_scale_encoder(sm_tokens, lg_tokens)
+
+        sm_cls, lg_cls = map(lambda t: t[:, 0], (sm_tokens, lg_tokens))
+
+        sm_logits = self.sm_mlp_head(sm_cls)
+        lg_logits = self.lg_mlp_head(lg_cls)
+
+        return sm_logits + lg_logits

vit_pytorch/deepvit.py

@@ -37,35 +37,41 @@ class Attention(nn.Module):
     def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
         super().__init__()
         inner_dim = dim_head *  heads
-        project_out = not (heads == 1 and dim_head == dim)
-
         self.heads = heads
         self.scale = dim_head ** -0.5
 
         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
 
+        self.reattn_weights = nn.Parameter(torch.randn(heads, heads))
+
+        self.reattn_norm = nn.Sequential(
+            Rearrange('b h i j -> b i j h'),
+            nn.LayerNorm(heads),
+            Rearrange('b i j h -> b h i j')
+        )
+
         self.to_out = nn.Sequential(
             nn.Linear(inner_dim, dim),
             nn.Dropout(dropout)
-        ) if project_out else nn.Identity()
+        )
 
-    def forward(self, x, mask = None):
+    def forward(self, x):
         b, n, _, h = *x.shape, self.heads
         qkv = self.to_qkv(x).chunk(3, dim = -1)
         q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
 
+        # attention
+
         dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
-        mask_value = -torch.finfo(dots.dtype).max
-
-        if mask is not None:
-            mask = F.pad(mask.flatten(1), (1, 0), value = True)
-            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
-            mask = rearrange(mask, 'b i -> b () i ()') * rearrange(mask, 'b j -> b () () j')
-            dots.masked_fill_(~mask, mask_value)
-            del mask
-
         attn = dots.softmax(dim=-1)
 
+        # re-attention
+
+        attn = einsum('b h i j, h g -> b g i j', attn, self.reattn_weights)
+        attn = self.reattn_norm(attn)
+
+        # aggregate and out
+
         out = einsum('b h i j, b h j d -> b h i d', attn, v)
         out = rearrange(out, 'b h n d -> b n (h d)')
         out =  self.to_out(out)
@@ -80,13 +86,13 @@ class Transformer(nn.Module):
                 Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
                 Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
             ]))
-    def forward(self, x, mask = None):
+    def forward(self, x):
         for attn, ff in self.layers:
-            x = attn(x, mask = mask)
+            x = attn(x)
             x = ff(x)
         return x
 
-class ViT(nn.Module):
+class DeepViT(nn.Module):
     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.):
         super().__init__()
         assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
@@ -113,7 +119,7 @@ class ViT(nn.Module):
             nn.Linear(dim, num_classes)
         )
 
-    def forward(self, img, mask = None):
+    def forward(self, img):
         x = self.to_patch_embedding(img)
         b, n, _ = x.shape
 
@@ -122,7 +128,7 @@ class ViT(nn.Module):
         x += self.pos_embedding[:, :(n + 1)]
         x = self.dropout(x)
 
-        x = self.transformer(x, mask)
+        x = self.transformer(x)
 
         x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
 

vit_pytorch/distill.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn.functional as F
 from torch import nn
-from vit_pytorch.vit_pytorch import ViT
+from vit_pytorch.vit import ViT
 from vit_pytorch.t2t import T2TViT
 from vit_pytorch.efficient import ViT as EfficientViT
 
@@ -15,7 +15,7 @@ def exists(val):
 # classes
 
 class DistillMixin:
-    def forward(self, img, distill_token = None, mask = None):
+    def forward(self, img, distill_token = None):
         distilling = exists(distill_token)
         x = self.to_patch_embedding(img)
         b, n, _ = x.shape
@@ -28,7 +28,7 @@ class DistillMixin:
             distill_tokens = repeat(distill_token, '() n d -> b n d', b = b)
             x = torch.cat((x, distill_tokens), dim = 1)
 
-        x = self._attend(x, mask)
+        x = self._attend(x)
 
         if distilling:
             x, distill_tokens = x[:, :-1], x[:, -1]
@@ -56,9 +56,9 @@ class DistillableViT(DistillMixin, ViT):
         v.load_state_dict(self.state_dict())
         return v
 
-    def _attend(self, x, mask):
+    def _attend(self, x):
         x = self.dropout(x)
-        x = self.transformer(x, mask)
+        x = self.transformer(x)
         return x
 
 class DistillableT2TViT(DistillMixin, T2TViT):
@@ -74,7 +74,7 @@ class DistillableT2TViT(DistillMixin, T2TViT):
         v.load_state_dict(self.state_dict())
         return v
 
-    def _attend(self, x, mask):
+    def _attend(self, x):
         x = self.dropout(x)
         x = self.transformer(x)
         return x
@@ -92,7 +92,7 @@ class DistillableEfficientViT(DistillMixin, EfficientViT):
         v.load_state_dict(self.state_dict())
         return v
 
-    def _attend(self, x, mask):
+    def _attend(self, x):
         return self.transformer(x)
 
 # knowledge distillation wrapper

vit_pytorch/mpp.py

@@ -0,0 +1,166 @@
+import math
+from functools import reduce
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from einops import rearrange, repeat
+
+# helpers
+
+
+def prob_mask_like(t, prob):
+    batch, seq_length, _ = t.shape
+    return torch.zeros((batch, seq_length)).float().uniform_(0, 1) < prob
+
+
+def get_mask_subset_with_prob(patched_input, prob):
+    batch, seq_len, _, device = *patched_input.shape, patched_input.device
+    max_masked = math.ceil(prob * seq_len)
+
+    rand = torch.rand((batch, seq_len), device=device)
+    _, sampled_indices = rand.topk(max_masked, dim=-1)
+
+    new_mask = torch.zeros((batch, seq_len), device=device)
+    new_mask.scatter_(1, sampled_indices, 1)
+    return new_mask.bool()
+
+
+# mpp loss
+
+
+class MPPLoss(nn.Module):
+    def __init__(self, patch_size, channels, output_channel_bits,
+                 max_pixel_val):
+        super(MPPLoss, self).__init__()
+        self.patch_size = patch_size
+        self.channels = channels
+        self.output_channel_bits = output_channel_bits
+        self.max_pixel_val = max_pixel_val
+
+    def forward(self, predicted_patches, target, mask):
+        # reshape target to patches
+        p = self.patch_size
+        target = rearrange(target,
+                           "b c (h p1) (w p2) -> b (h w) c (p1 p2) ",
+                           p1=p,
+                           p2=p)
+
+        avg_target = target.mean(dim=3)
+
+        bin_size = self.max_pixel_val / self.output_channel_bits
+        channel_bins = torch.arange(bin_size, self.max_pixel_val, bin_size)
+        discretized_target = torch.bucketize(avg_target, channel_bins)
+        discretized_target = F.one_hot(discretized_target,
+                                       self.output_channel_bits)
+        c, bi = self.channels, self.output_channel_bits
+        discretized_target = rearrange(discretized_target,
+                                       "b n c bi -> b n (c bi)",
+                                       c=c,
+                                       bi=bi)
+
+        bin_mask = 2**torch.arange(c * bi - 1, -1,
+                                   -1).to(discretized_target.device,
+                                          discretized_target.dtype)
+        target_label = torch.sum(bin_mask * discretized_target, -1)
+
+        predicted_patches = predicted_patches[mask]
+        target_label = target_label[mask]
+        loss = F.cross_entropy(predicted_patches, target_label)
+        return loss
+
+
+# main class
+
+
+class MPP(nn.Module):
+    def __init__(self,
+                 transformer,
+                 patch_size,
+                 dim,
+                 output_channel_bits=3,
+                 channels=3,
+                 max_pixel_val=1.0,
+                 mask_prob=0.15,
+                 replace_prob=0.5,
+                 random_patch_prob=0.5):
+        super().__init__()
+
+        self.transformer = transformer
+        self.loss = MPPLoss(patch_size, channels, output_channel_bits,
+                            max_pixel_val)
+
+        # output transformation
+        self.to_bits = nn.Linear(dim, 2**(output_channel_bits * channels))
+
+        # vit related dimensions
+        self.patch_size = patch_size
+
+        # mpp related probabilities
+        self.mask_prob = mask_prob
+        self.replace_prob = replace_prob
+        self.random_patch_prob = random_patch_prob
+
+        # token ids
+        self.mask_token = nn.Parameter(torch.randn(1, 1, dim * channels))
+
+    def forward(self, input, **kwargs):
+        transformer = self.transformer
+        # clone original image for loss
+        img = input.clone().detach()
+
+        # reshape raw image to patches
+        p = self.patch_size
+        input = rearrange(input,
+                          'b c (h p1) (w p2) -> b (h w) (p1 p2 c)',
+                          p1=p,
+                          p2=p)
+
+        mask = get_mask_subset_with_prob(input, self.mask_prob)
+
+        # mask input with mask patches with probability of `replace_prob` (keep patches the same with probability 1 - replace_prob)
+        masked_input = input.clone().detach()
+
+        # if random token probability > 0 for mpp
+        if self.random_patch_prob > 0:
+            random_patch_sampling_prob = self.random_patch_prob / (
+                1 - self.replace_prob)
+            random_patch_prob = prob_mask_like(input,
+                                               random_patch_sampling_prob)
+            bool_random_patch_prob = mask * random_patch_prob == True
+            random_patches = torch.randint(0,
+                                           input.shape[1],
+                                           (input.shape[0], input.shape[1]),
+                                           device=input.device)
+            randomized_input = masked_input[
+                torch.arange(masked_input.shape[0]).unsqueeze(-1),
+                random_patches]
+            masked_input[bool_random_patch_prob] = randomized_input[
+                bool_random_patch_prob]
+
+        # [mask] input
+        replace_prob = prob_mask_like(input, self.replace_prob)
+        bool_mask_replace = (mask * replace_prob) == True
+        masked_input[bool_mask_replace] = self.mask_token
+
+        # linear embedding of patches
+        masked_input = transformer.to_patch_embedding[-1](masked_input)
+
+        # add cls token to input sequence
+        b, n, _ = masked_input.shape
+        cls_tokens = repeat(transformer.cls_token, '() n d -> b n d', b=b)
+        masked_input = torch.cat((cls_tokens, masked_input), dim=1)
+
+        # add positional embeddings to input
+        masked_input += transformer.pos_embedding[:, :(n + 1)]
+        masked_input = transformer.dropout(masked_input)
+
+        # get generator output and get mpp loss
+        masked_input = transformer.transformer(masked_input, **kwargs)
+        cls_logits = self.to_bits(masked_input)
+        logits = cls_logits[:, 1:, :]
+
+        mpp_loss = self.loss(logits, img, mask)
+
+        return mpp_loss

vit_pytorch/recorder.py

@@ -0,0 +1,54 @@
+from functools import wraps
+import torch
+from torch import nn
+
+from vit_pytorch.vit import Attention
+
+def find_modules(nn_module, type):
+    return [module for module in nn_module.modules() if isinstance(module, type)]
+
+class Recorder(nn.Module):
+    def __init__(self, vit):
+        super().__init__()
+        self.vit = vit
+
+        self.data = None
+        self.recordings = []
+        self.hooks = []
+        self.hook_registered = False
+        self.ejected = False
+
+    def _hook(self, _, input, output):
+        self.recordings.append(output.clone().detach())
+
+    def _register_hook(self):
+        modules = find_modules(self.vit.transformer, Attention)
+        for module in modules:
+            handle = module.attend.register_forward_hook(self._hook)
+            self.hooks.append(handle)
+        self.hook_registered = True
+
+    def eject(self):
+        self.ejected = True
+        for hook in self.hooks:
+            hook.remove()
+        self.hooks.clear()
+        return self.vit
+
+    def clear(self):
+        self.recordings.clear()
+
+    def record(self, attn):
+        recording = attn.clone().detach()
+        self.recordings.append(recording)
+
+    def forward(self, img):
+        assert not self.ejected, 'recorder has been ejected, cannot be used anymore'
+        self.clear()
+
+        if not self.hook_registered:
+            self._register_hook()
+
+        pred = self.vit(img)
+        attns = torch.stack(self.recordings, dim = 1)
+        return pred, attns

vit_pytorch/t2t.py

@@ -2,16 +2,21 @@ import math
 import torch
 from torch import nn
 
-from vit_pytorch.vit_pytorch import Transformer
+from vit_pytorch.vit import Transformer
 
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
-# classes
+# helpers
+
+def exists(val):
+    return val is not None
 
 def conv_output_size(image_size, kernel_size, stride, padding):
     return int(((image_size - kernel_size + (2 * padding)) / stride) + 1)
 
+# classes
+
 class RearrangeImage(nn.Module):
     def forward(self, x):
         return rearrange(x, 'b (h w) c -> b c h w', h = int(math.sqrt(x.shape[1])))
@@ -19,8 +24,7 @@ class RearrangeImage(nn.Module):
 # main class
 
 class T2TViT(nn.Module):
-    def __init__(
-        self, *, image_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., t2t_layers = ((7, 4), (3, 2), (3, 2))):
+    def __init__(self, *, image_size, num_classes, dim, depth = None, heads = None, mlp_dim = None, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0., transformer = None, t2t_layers = ((7, 4), (3, 2), (3, 2))):
         super().__init__()
         assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
 
@@ -47,7 +51,11 @@ class T2TViT(nn.Module):
         self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
         self.dropout = nn.Dropout(emb_dropout)
 
-        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+        if not exists(transformer):
+            assert all([exists(depth), exists(heads), exists(mlp_dim)]), 'depth, heads, and mlp_dim must be supplied'
+            self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+        else:
+            self.transformer = transformer
 
         self.pool = pool
         self.to_latent = nn.Identity()

vit_pytorch/vit.py

@@ -0,0 +1,115 @@
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.fn = fn
+    def forward(self, x, **kwargs):
+        return self.fn(self.norm(x), **kwargs)
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.attend = nn.Softmax(dim = -1)
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        b, n, _, h = *x.shape, self.heads
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)
+
+        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+
+        attn = self.attend(dots)
+
+        out = einsum('b h i j, b h j d -> b h i d', attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
+        super().__init__()
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
+                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
+            ]))
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return x
+
+class ViT(nn.Module):
+    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.):
+        super().__init__()
+        assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
+        num_patches = (image_size // patch_size) ** 2
+        patch_dim = channels * patch_size ** 2
+        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
+
+        self.to_patch_embedding = nn.Sequential(
+            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
+            nn.Linear(patch_dim, dim),
+        )
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
+
+        self.pool = pool
+        self.to_latent = nn.Identity()
+
+        self.mlp_head = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, num_classes)
+        )
+
+    def forward(self, img):
+        x = self.to_patch_embedding(img)
+        b, n, _ = x.shape
+
+        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x += self.pos_embedding[:, :(n + 1)]
+        x = self.dropout(x)
+
+        x = self.transformer(x)
+
+        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return self.mlp_head(x)