use einops

.isort.cfg

@@ -1,2 +1,2 @@
 [settings]
-known_third_party = PIL,matplotlib,numpy,pytest,pytorch_lightning,torch,torchmetrics,torchvision
+known_third_party = PIL,einops,matplotlib,numpy,pytest,pytorch_lightning,torch,torchmetrics,torchvision

pyproject.toml

@@ -8,6 +8,7 @@ authors = ["sami jaghouar <[email protected]>"]
 python = "^3.8"
 torch = "^1.9.1"
 torchvision = "^0.10.1"
+einops = "^0.3.2"
 
 [tool.poetry.dev-dependencies]
 pytest = "^5.2"

vision_transformers/vit.py

@@ -1,8 +1,8 @@
-from itertools import product
 from typing import Tuple
 
 import torch
 import torch.nn as nn
+from einops.layers.torch import Rearrange
 
 
 class PatchSizeError(ValueError):
@@ -24,14 +24,9 @@ def __init__(self, size_of_patch: int, input_shape: Tuple[int, int]):
         self.input_shape = input_shape
         self.N = (input_shape[0] // self.P) * (input_shape[1] // self.P)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor():
-        return torch.stack(
-            [
-                x[:, :, i * self.P : (i + 1) * self.P, j * self.P : (j + 1) * self.P]
-                for i, j in product(
-                    range(0, self.input_shape[0] // self.P),
-                    range(0, self.input_shape[1] // self.P),
-                )
-            ],
-            dim=1,
+        self.rearrange = Rearrange(
+            " n c (h p1) (w p2) -> n (h w) c p1 p2", p2=self.P, p1=self.P
         )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor():
+        return self.rearrange(x)