remove patch size from T2TViT

setup.py

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

vit_pytorch/distill.py

@@ -2,7 +2,6 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 from vit_pytorch.vit_pytorch import ViT
-from vit_pytorch.t2t import T2TViT
 from vit_pytorch.efficient import ViT as EfficientViT
 
 from einops import rearrange, repeat
@@ -61,24 +60,6 @@ class DistillableViT(DistillMixin, ViT):
         x = self.transformer(x, mask)
         return x
 
-class DistillableT2TViT(DistillMixin, T2TViT):
-    def __init__(self, *args, **kwargs):
-        super(DistillableT2TViT, self).__init__(*args, **kwargs)
-        self.args = args
-        self.kwargs = kwargs
-        self.dim = kwargs['dim']
-        self.num_classes = kwargs['num_classes']
-
-    def to_vit(self):
-        v = T2TViT(*self.args, **self.kwargs)
-        v.load_state_dict(self.state_dict())
-        return v
-
-    def _attend(self, x, mask):
-        x = self.dropout(x)
-        x = self.transformer(x)
-        return x
-
 class DistillableEfficientViT(DistillMixin, EfficientViT):
     def __init__(self, *args, **kwargs):
         super(DistillableEfficientViT, self).__init__(*args, **kwargs)
@@ -107,7 +88,7 @@ class DistillWrapper(nn.Module):
         alpha = 0.5
     ):
         super().__init__()
-        assert (isinstance(student, (DistillableViT, DistillableT2TViT, DistillableEfficientViT))) , 'student must be a vision transformer'
+        assert (isinstance(student, (DistillableViT, DistillableEfficientViT))) , 'student must be a vision transformer'
 
         self.teacher = teacher
         self.student = student

vit_pytorch/t2t.py

@@ -7,16 +7,11 @@ from vit_pytorch.vit_pytorch import Transformer
 from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
-# helpers
-
-def exists(val):
-    return val is not None
+# classes
 
 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])))
@@ -25,7 +20,7 @@ class RearrangeImage(nn.Module):
 
 class T2TViT(nn.Module):
     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))):
+        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))):
         super().__init__()
         assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
 
@@ -52,11 +47,7 @@ class T2TViT(nn.Module):
         self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
         self.dropout = nn.Dropout(emb_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.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
 
         self.pool = pool
         self.to_latent = nn.Identity()
@@ -72,7 +63,7 @@ class T2TViT(nn.Module):
 
         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
+        x += self.pos_embedding[:, :(n + 1)]
         x = self.dropout(x)
 
         x = self.transformer(x)

vit_pytorch/vit_pytorch.py

@@ -1,3 +1,4 @@
+import math
 import torch
 from torch import nn, einsum
 import torch.nn.functional as F
@@ -37,17 +38,14 @@ 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.to_out = nn.Sequential(
             nn.Linear(inner_dim, dim),
             nn.Dropout(dropout)
-        ) if project_out else nn.Identity()
+        )
 
     def forward(self, x, mask = None):
         b, n, _, h = *x.shape, self.heads