Fix generate caption function

source/predict_sample.py

@@ -14,7 +14,10 @@ def generate_caption(image: torch.Tensor,
                      image_decoder: Decoder,
                      vocab: Vocab,
                      device: torch.device) -> list[str]:
-    """ Generate caption of a single image of size (1, 3, 224, 224)
+    """
+    Generate caption of a single image of size (3, 224, 224).
+    Generating of caption starts with <sos>, and each next predicted word ID
+    is appended for the next LSTM input until the sentence reaches MAX_LENGTH or <eos>.
 
     Returns:
         list[str]: caption for given image
@@ -25,49 +28,57 @@ def generate_caption(image: torch.Tensor,
     image = image.unsqueeze(0)
     # image: (1, 3, 224, 224)
 
-    features = image_encoder.forward(image)
-    # features: (1, IMAGE_EMB_DIM)
-    features = features.to(device)
-    features = features.unsqueeze(0)
-    # features: (1, 1, IMAGE_EMB_DIM)
-
     hidden = image_decoder.hidden_state_0
     cell = image_decoder.cell_state_0
     # hidden, cell : (NUM_LAYER, 1, HIDDEN_DIM)
 
     sentence = []
 
-    # start with '<sos>' as first word
-    previous_word = vocab.index2word[vocab.SOS]
+    # initialize LSTM input to SOS token = 1
+    input_words = [vocab.SOS]
 
     MAX_LENGTH = 20
 
     for i in range(MAX_LENGTH):
 
-        input_word_id = vocab.word_to_index(previous_word)
-        input_word_tensor = torch.tensor([input_word_id]).unsqueeze(0)
-        # input_word_tensor : (1, 1)
+        features = image_encoder.forward(image)
+        # features: (1, IMAGE_EMB_DIM)
+        features = features.to(device)
+        features = features.unsqueeze(0)
+        # features: (1, 1, IMAGE_EMB_DIM)
+
+        input_words_tensor = torch.tensor([input_words])
+        # input_word_tensor : (B=1, SEQ_LENGTH)
+        input_words_tensor = input_words_tensor.to(device)
+
+        lstm_input = emb_layer.forward(input_words_tensor)
+        # lstm_input : (B=1, SEQ_LENGTH, WORD_EMB_DIM)
+
+        lstm_input = lstm_input.permute(1, 0, 2)
+        # lstm_input : (SEQ_LENGTH, B=1, WORD_EMB_DIM)
+        SEQ_LENGTH = lstm_input.shape[0]
 
-        input_word_tensor = input_word_tensor.to(device)
-        lstm_input = emb_layer.forward(input_word_tensor)
-        # lstm_input : (1, 1, WORD_EMB_DIM)
+        features = features.repeat(SEQ_LENGTH, 1, 1)
+        # features : (SEQ_LENGTH, B=1, IMAGE_EMB_DIM)
 
-        next_word_pred, (hidden, cell) = image_decoder.forward(lstm_input, features, hidden, cell)
-        # next_word_pred : (1, 1, VOCAB_SIZE)
+        next_id_pred, (hidden, cell) = image_decoder.forward(lstm_input, features, hidden, cell)
+        # next_id_pred : (SEQ_LENGTH, 1, VOCAB_SIZE)
 
-        next_word_pred = next_word_pred[0, 0, :]
-        # next_word_pred : (VOCAB_SIZE)
+        next_id_pred = next_id_pred[-1, 0, :]
+        # next_id_pred : (VOCAB_SIZE)
+        next_id_pred = torch.argmax(next_id_pred)
 
-        next_word_pred = torch.argmax(next_word_pred)
-        next_word_pred = vocab.index_to_word(int(next_word_pred.item()))
+        # append it to input_words which will be again as input for LSTM
+        input_words.append(next_id_pred.item())
+
+        # id --> word
+        next_word_pred = vocab.index_to_word(int(next_id_pred.item()))
+        sentence.append(next_word_pred)
 
         # stop if we predict '<eos>'
         if next_word_pred == vocab.index2word[vocab.EOS]:
             break
 
-        sentence.append(next_word_pred)
-        previous_word = next_word_pred
-
     return sentence