cleanup

vit_pytorch/distill.py

@@ -43,6 +43,36 @@ class DistillableViT(ViT):
         x = self.to_latent(x)
         return self.mlp_head(x), distill_tokens
 
+class DistillableEfficientViT(EfficientViT):
+    def __init__(self, *args, **kwargs):
+        super(DistillableEfficientViT, self).__init__(*args, **kwargs)
+        self.dim = kwargs['dim']
+        self.num_classes = kwargs['num_classes']
+
+    def forward(self, img, distill_token, mask = None):
+        p = self.patch_size
+
+        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)
+        x = self.patch_to_embedding(x)
+        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)]
+
+        distill_tokens = repeat(distill_token, '() n d -> b n d', b = b)
+        x = torch.cat((x, distill_tokens), dim = 1)
+
+        x = self.transformer(x)
+
+        x, distill_tokens = x[:, :-1], x[:, -1]
+        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
+
+        x = self.to_latent(x)
+        return self.mlp_head(x), distill_tokens
+
+# knowledge distillation wrapper
+
 class DistillWrapper(nn.Module):
     def __init__(
         self,
@@ -53,11 +83,7 @@ class DistillWrapper(nn.Module):
         alpha = 0.5
     ):
         super().__init__()
-        assert (
-                isinstance(student, DistillableViT)
-                or
-                isinstance(student, DistillableEfficientViT)
-        ) , 'student must be a vision transformer'
+        assert (isinstance(student, (DistillableViT, DistillableEfficientViT))) , 'student must be a vision transformer'
 
         self.teacher = teacher
         self.student = student
@@ -95,33 +121,3 @@ class DistillWrapper(nn.Module):
         distill_loss *= T ** 2
 
         return loss * alpha + distill_loss * (1 - alpha)
-
-
-class DistillableEfficientViT(EfficientViT):
-    def __init__(self, *args, **kwargs):
-        super(DistillableEfficientViT, self).__init__(*args, **kwargs)
-        self.dim = kwargs['dim']
-        self.num_classes = kwargs['num_classes']
-
-    def forward(self, img, distill_token, mask = None):
-        p = self.patch_size
-
-        x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p, p2 = p)
-        x = self.patch_to_embedding(x)
-        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)]
-
-        distill_tokens = repeat(distill_token, '() n d -> b n d', b = b)
-        x = torch.cat((x, distill_tokens), dim = 1)
-
-        x = self.transformer(x)
-
-        x, distill_tokens = x[:, :-1], x[:, -1]
-        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
-
-        x = self.to_latent(x)
-        return self.mlp_head(x), distill_tokens
-