Update code and weights

source/config.py

@@ -12,8 +12,8 @@ class Config(object):
         self.VOCAB_SIZE = 5000
 
         self.NUM_LAYER = 1
-        self.IMAGE_EMB_DIM = 256
-        self.WORD_EMB_DIM = 256
+        self.IMAGE_EMB_DIM = 512
+        self.WORD_EMB_DIM = 5121
         self.HIDDEN_DIM = 512
 
         self.EMBEDDING_WEIGHT_FILE = 'source/weights/embeddings-32B-512H-1L-e5.pt'

source/model.py

@@ -1,124 +1,117 @@
 import torch
-import torch._utils
 import torch.nn as nn
 import torchvision.models as models
 from typing import Tuple
-from source.config import Config
 
 
 class Encoder(nn.Module):
-    def __init__(self, image_emb_dim: int, device: torch.device):
-        """ Image encoder to obtain features from images. Contains pretrained Resnet50 with last layer removed
-            and a linear layer with the output dimension of (BATCH, image_emb_dim)
-        """
+    """
+    Image encoder to obtain features from images using a pretrained ResNet-50 model.
+    The last layer of ResNet-50 is removed, and a linear layer is added to transform
+    the output to the desired feature dimension.
+
+    Args:
+        image_emb_dim (int): Final output dimension of image features.
+        device (torch.device): Device to run the model on (CPU or GPU).
+    """
 
+    def __init__(self, image_emb_dim: int, device: torch.device):
         super(Encoder, self).__init__()
         self.image_emb_dim = image_emb_dim
         self.device = device
 
-        # pretrained Resnet50 model with freezed parameters
+        # Load pretrained ResNet-50 model and freeze its parameters
         resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
         for param in resnet.parameters():
             param.requires_grad_(False)
 
-        # remove last layer
+        # Remove the last layer of ResNet-50
         modules = list(resnet.children())[:-1]
         self.resnet = nn.Sequential(*modules)
 
-        # define a final classifier
-        self.fc = nn.Linear(in_features=resnet.fc.in_features, out_features=self.image_emb_dim)
+        # Define a final classifier
+        self.fc = nn.Linear(resnet.fc.in_features, self.image_emb_dim)
 
     def forward(self, images: torch.Tensor) -> torch.Tensor:
-        """ Forward operation of encoder, passing images through resnet and then linear layer.
+        """
+        Forward pass through the encoder.
 
         Args:
-            > images (torch.Tensor): (BATCH, 3, 224, 224)
+            images (torch.Tensor): Input images of shape (BATCH, 3, 224, 224).
 
         Returns:
-            > features (torch.Tensor): (BATCH, IMAGE_EMB_DIM)
+            torch.Tensor: Image features of shape (BATCH, IMAGE_EMB_DIM).
         """
-
         features = self.resnet(images)
-        # features: (BATCH, 2048, 1, 1)
-
+        # Reshape features to (BATCH, 2048)
         features = features.reshape(features.size(0), -1).to(self.device)
-        # features: (BATCH, 2048)
-
+        # Pass features through final linear layer
         features = self.fc(features).to(self.device)
-        # features: (BATCH, IMAGE_EMB_DIM)
-
         return features
 
 
 class Decoder(nn.Module):
+    """
+    Decoder that uses an LSTM to generate captions from embedded words and encoded image features.
+    The hidden and cell states of the LSTM are initialized using the encoded image features.
+
+    Args:
+        word_emb_dim (int): Dimension of word embeddings.
+        hidden_dim (int): Dimension of the LSTM hidden state.
+        num_layers (int): Number of LSTM layers.
+        vocab_size (int): Size of the vocabulary (output dimension of the final linear layer).
+        device (torch.device): Device to run the model on (CPU or GPU).
+    """
+
     def __init__(self,
-                 image_emb_dim: int,
                  word_emb_dim: int,
                  hidden_dim: int,
                  num_layers: int,
                  vocab_size: int,
                  device: torch.device):
-        """
-        Decoder taking as input for the LSTM layer the concatenation of features obtained from the encoder
-        and embedded captions obtained from the embedding layer. Hidden and cell states are randomly initialized.
-        Final classifier is a linear layer with output dimension of the size of a vocabulary.
-        """
-
         super(Decoder, self).__init__()
 
-        self.config = Config()
-
-        self.image_emd_dim = image_emb_dim
         self.word_emb_dim = word_emb_dim
         self.hidden_dim = hidden_dim
-        self.num_layer = num_layers
+        self.num_layers = num_layers
         self.vocab_size = vocab_size
         self.device = device
 
-        self.hidden_state_0 = nn.Parameter(torch.zeros((self.num_layer, 1, self.hidden_dim)))
-        self.cell_state_0 = nn.Parameter(torch.zeros((self.num_layer, 1, self.hidden_dim)))
+        # Initialize hidden and cell states
+        self.hidden_state_0 = nn.Parameter(torch.zeros((self.num_layers, 1, self.hidden_dim)))
+        self.cell_state_0 = nn.Parameter(torch.zeros((self.num_layers, 1, self.hidden_dim)))
 
-        self.lstm = nn.LSTM(input_size=self.image_emd_dim + self.word_emb_dim,
-                            hidden_size=self.hidden_dim,
-                            num_layers=self.num_layer,
+        # Define LSTM layer
+        self.lstm = nn.LSTM(self.word_emb_dim,
+                            self.hidden_dim,
+                            num_layers=self.num_layers,
                             bidirectional=False)
 
+        # Define final linear layer with LogSoftmax activation
         self.fc = nn.Sequential(
-            nn.Linear(in_features=self.hidden_dim, out_features=self.vocab_size),
+            nn.Linear(self.hidden_dim, self.vocab_size),
             nn.LogSoftmax(dim=2)
         )
 
     def forward(self,
                 embedded_captions: torch.Tensor,
-                features: torch.Tensor,
                 hidden: torch.Tensor,
                 cell: torch.Tensor) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         """
-        Forward operation of (word-by-word) decoder.
-        The LSTM input (concatenation of embedded_captions and features) is passed through LSTM and then linear layer.
+        Forward pass through the decoder.
 
         Args:
-
-            > embedded_captions(torch.Tensor): (SEQ_LENGTH = 1, BATCH, WORD_EMB_DIM)
-            > features (torch.Tensor): (1, BATCH, IMAGE_EMB_DIM)
-            > hidden (torch.Tensor): (NUM_LAYER, BATCH, HIDDEN_DIM)
-            > cell (torch.Tensor): (NUM_LAYER, BATCH, HIDDEN_DIM)
+            embedded_captions (torch.Tensor): Embedded captions of shape (SEQ_LEN, BATCH, WORD_EMB_DIM).
+            hidden (torch.Tensor): LSTM hidden state of shape (NUM_LAYER, BATCH, HIDDEN_DIM).
+            cell (torch.Tensor): LSTM cell state of shape (NUM_LAYER, BATCH, HIDDEN_DIM).
 
         Returns:
-
-            > output (torch.Tensor): (1, BATCH, VOCAB_SIZE)
-            > (hidden, cell) (torch.Tensor, torch.Tensor): (NUM_LAYER, BATCH, HIDDEN_DIM)
+            Tuple:
+                - output (torch.Tensor): Output logits of shape (SEQ_LEN, BATCH, VOCAB_SIZE).
+                - (hidden, cell) (Tuple[torch.Tensor, torch.Tensor]): Updated hidden and cell states.
         """
-
-        lstm_input = torch.cat((embedded_captions, features), dim=2)
-
-        output, (hidden, cell) = self.lstm(lstm_input, (hidden, cell))
-        # output : (SEQ_LENGTH, BATCH, HIDDEN_DIM)
-        # hidden : (NUM_LAYER, BATCH, HIDDEN_DIM)
-
-        output = output.to(self.device)
-
+        # Pass through LSTM
+        output, (hidden, cell) = self.lstm(embedded_captions, (hidden, cell))
+        # Pass through final linear layer
         output = self.fc(output)
-        # output : (SEQ_LENGTH, BATCH, VOCAB_SIZE)
-
         return output, (hidden, cell)

source/predict_sample.py

@@ -104,8 +104,7 @@ def main_caption(image):
     emb_layer = torch.nn.Embedding(num_embeddings=config.VOCAB_SIZE,
                                    embedding_dim=config.WORD_EMB_DIM,
                                    padding_idx=vocab.PADDING_INDEX)
-    image_decoder = Decoder(image_emb_dim=config.IMAGE_EMB_DIM,
-                            word_emb_dim=config.WORD_EMB_DIM,
+    image_decoder = Decoder(word_emb_dim=config.WORD_EMB_DIM,
                             hidden_dim=config.HIDDEN_DIM,
                             num_layers=config.NUM_LAYER,
                             vocab_size=config.VOCAB_SIZE,

source/weights/{decoder-32B-512H-1L-e2.pt → decoder-32B-512H-1L-e10.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c14f313a1fea17eb147567a456f418355e666858a0f0fa4f5dfa8f8a561e076a
-size 18671955
+oid sha256:435a74d3029be0e1bce2dd451cbb58ec84a2e9ee2e3d685fd9e151c5a2123139
+size 18671964

source/weights/decoder-32B-512H-1L-e6.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:1914f17b249f0819e2680740e1bed990e38cde2fd5db916e3f33b2e106f6c2fc
-size 18671955

source/weights/{decoder-32B-512H-1L-e4.pt → embeddings-32B-512H-1L-e10.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:92a9289a063a101f4f3214cc7b67990d62b9054dfe917cb40492a7bde5440c60
-size 18671955
+oid sha256:1e6b0a7b05ab93d06da4fcc93dff769d02fc3ff48963b6979d3faa00de6f62a9
+size 10241467

source/weights/embeddings-32B-512H-1L-e2.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:efb8031895da44da642975ba1a1997a214437ca61113edbbfa31f30a26c2ad9e
-size 5121462

source/weights/embeddings-32B-512H-1L-e4.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9a0faba2400080ae7acf50c38b214f389a763c95f2f587d1d664110b5d9978cf
-size 5121462

source/weights/embeddings-32B-512H-1L-e5.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3ad73c03e1547417874d7d154213a893ac38adb24d74386a2055fc4d1fd46884
-size 5121041

source/weights/embeddings-32B-512H-1L-e6.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:76efafb94073fa60b15cfce698f78072465067e428968a104d174b8a3adabd32
-size 5121462

source/weights/{decoder-32B-512H-1L-e5.pt → encoder-32B-512H-1L-e10.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5dd4bb9bd858cc8518a4af612df8721ca67d40a3428f53d34c50baef4ee87371
-size 18671739
+oid sha256:ea482f42ec88705fef214bfa92acd4ee535e331110eaeda32198e63a8a9c108c
+size 98552306

source/weights/encoder-32B-512H-1L-e2.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0dff5de9d9ad9ea43fc5f67798b610f0bb92224590eba264766921b418a0d7a6
-size 96453806

source/weights/encoder-32B-512H-1L-e4.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c51a0ffb6eccb3fc2163b7c3214bdb9e32972a14b12d2be210289865bec4d7f7
-size 96453806

source/weights/encoder-32B-512H-1L-e5.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:92c5362bce19b36b330c58455985136d546821404d31477947544af70dbeab83
-size 96458817

source/weights/encoder-32B-512H-1L-e6.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:9d59e51481084ee51bd810e7b0b87fa89577cfdcc8cfd76d5495f45beaff9feb
-size 96453806