image embedding and encoding

model.py

@@ -0,0 +1,57 @@
+from einops.layers.torch import Rearrange
+import einops
+import math
+import torch.nn as nn
+import torch
+
+
+class ImageEmbedding(nn.Module):
+    """Reshape image into patches and project into given dimension"""
+
+    def __init__(self, d_model, input_height, input_width, patch_size=16):
+        super().__init__()
+        assert input_height % patch_size == 0 and input_width % patch_size == 0, \
+            "Cannot split image in patches"
+
+        self.tokenize = Rearrange(
+            'b c (h1 h2) (w1 w2) -> b (c h1 w1) (h2 w2)',
+            h2=patch_size,
+            w2=patch_size
+        )
+        self.projection = nn.Linear(patch_size ** 2, d_model)
+
+    def forward(self, image_batch):
+        image_batch = self.tokenize(image_batch)
+        image_batch = self.projection(image_batch)
+        return image_batch
+
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, max_sequence_len=5000):
+        super().__init__()
+
+        # pos - position in sequence, i - index of element embedding
+        # PE(pos, 2i) = sin(pos / 10000**(2i / d_model)) = sin(pos * e**(2i * (-log(10000))/d_model))
+        # PE(pos, 2i+1) = cos(pos / 10000**(2i / d_model)) = cos(pos * e**(2i * (-log(10000))/d_model))
+
+        positions = torch.arange(max_sequence_len)
+        even_embedding_indices = torch.arange(0, d_model, 2)
+
+        expression = torch.exp(even_embedding_indices * (-math.log(10000.0) / d_model))
+        expression = torch.einsum("i, j -> ij", positions, expression)
+
+        even_encodings = torch.sin(expression)
+        odd_encodings = torch.cos(expression)
+
+        positional_encodings = einops.rearrange(
+            [even_encodings, odd_encodings],
+            'even_odd pos embed -> pos 1 (embed even_odd)'
+        )
+
+        self.register_buffer('positional_encodings', positional_encodings)
+
+    def forward(self, image_batch):
+        batch_size = image_batch.size(0)
+        positional_encodings = self.positional_encodings[:batch_size]
+        return positional_encodings