Create trainning.py

trainning.py

@@ -0,0 +1,301 @@
+import torch
+import string
+import torch.nn.functional as F
+import torch.nn as nn
+import torchvision.models as models
+
+def decoder(indices, vocab):
+    tokens = [vocab.lookup_token(idx) for idx in indices]
+    words = []
+    current_word = []
+    for token in tokens:
+        if len(token) == 1 and token in string.ascii_lowercase:
+            current_word.append(token)
+        else:
+            if current_word:
+                words.append("".join(current_word))
+                current_word = []
+            words.append(token)
+
+    if current_word:
+        words.append(" "+"".join(current_word))
+
+    return "".join(words)
+
+def beam_search_caption(model, images, vocab, decoder, device="cpu",
+                       start_token="<sos>", end_token="<eos>",
+                       beam_width=3, max_seq_length=100):
+    """
+    Generates captions for images using beam search.
+
+    Args:
+        model (ImgCap): The image captioning model.
+        images (torch.Tensor): Batch of images.
+        vocab (Vocab): Vocabulary object.
+        decoder (function): Function to decode indices to words.
+        device (str): Device to perform computation on.
+        start_token (str): Start-of-sequence token.
+        end_token (str): End-of-sequence token.
+        beam_width (int): Number of beams to keep.
+        max_seq_length (int): Maximum length of the generated caption.
+
+    Returns:
+        list: Generated captions for each image in the batch.
+    """
+    model.eval()
+
+    with torch.no_grad():
+        start_index = vocab[start_token]
+        end_index = vocab[end_token]
+        images = images.to(device)
+        batch_size = images.size(0)
+        
+        # Ensure batch_size is 1 for beam search (one image at a time)
+        if batch_size != 1:
+            raise ValueError("Beam search currently supports batch_size=1.")
+
+        cnn_feature = model.cnn(images)  # Shape: (1, 1024)
+        lstm_input = model.lstm.projection(cnn_feature).unsqueeze(1)  # Shape: (1, 1, 1024)
+        state = None  # Initial LSTM state
+
+        # Initialize the beam with the start token
+        sequences = [([start_index], 0.0, lstm_input, state)]  # List of tuples: (sequence, score, input, state)
+
+        completed_sequences = []
+
+        for _ in range(max_seq_length):
+            all_candidates = []
+
+            # Iterate over all current sequences in the beam
+            for seq, score, lstm_input, state in sequences:
+                # If the last token is the end token, add the sequence to completed_sequences
+                if seq[-1] == end_index:
+                    completed_sequences.append((seq, score))
+                    continue
+
+                # Pass the current input and state through the LSTM
+                lstm_out, state_new = model.lstm.lstm(lstm_input, state)  # lstm_out: (1, 1, 1024)
+
+                # Pass the LSTM output through the fully connected layer to get logits
+                output = model.lstm.fc(lstm_out.squeeze(1))  # Shape: (1, vocab_size)
+
+                # Compute log probabilities
+                log_probs = F.log_softmax(output, dim=1)  # Shape: (1, vocab_size)
+
+                # Get the top beam_width tokens and their log probabilities
+                top_log_probs, top_indices = log_probs.topk(beam_width, dim=1)  # Each of shape: (1, beam_width)
+
+                # Iterate over the top tokens to create new candidate sequences
+                for i in range(beam_width):
+                    token = top_indices[0, i].item()
+                    token_log_prob = top_log_probs[0, i].item()
+
+                    # Create a new sequence by appending the current token
+                    new_seq = seq + [token]
+                    new_score = score + token_log_prob
+
+                    # Get the embedding of the new token
+                    token_tensor = torch.tensor([token], device=device)
+                    new_lstm_input = model.lstm.embedding(token_tensor).unsqueeze(1)  # Shape: (1, 1, 1024)
+
+                    # Clone the new state to ensure each beam has its own state
+                    if state_new is not None:
+                        new_state = (state_new[0].clone(), state_new[1].clone())
+                    else:
+                        new_state = None
+
+                    # Add the new candidate to all_candidates
+                    all_candidates.append((new_seq, new_score, new_lstm_input, new_state))
+
+            # If no candidates are left to process, break out of the loop
+            if not all_candidates:
+                break
+
+            # Sort all candidates by score in descending order
+            ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True)
+
+            # Select the top beam_width sequences to form the new beam
+            sequences = ordered[:beam_width]
+
+            # If enough completed sequences are found, stop early
+            if len(completed_sequences) >= beam_width:
+                break
+
+        # If no sequences have completed, use the current sequences
+        if len(completed_sequences) == 0:
+            completed_sequences = sequences
+
+        # Select the sequence with the highest score
+        best_seq, best_score = max(completed_sequences, key=lambda x: x[1])
+
+        if best_seq[0] == start_index:
+            best_seq = best_seq[1:]
+
+        best_caption = decoder(best_seq, vocab)
+
+    return best_caption
+
+
+def generate_caption(model, images, vocab, decoder, device="cpu", start_token="<sos>", end_token="<eos>", max_seq_length=100, top_k=2):
+    model.eval()
+
+    with torch.no_grad():
+        start_index = vocab[start_token]
+        end_index = vocab[end_token]
+        images = images.to(device)
+        batch_size = images.size(0)
+
+        end_token_appear = {i: False for i in range(batch_size)}
+        captions = [[] for _ in range(batch_size)]
+
+        cnn_feature = model.cnn(images)
+        lstm_input = model.lstm.projection(cnn_feature).unsqueeze(1)  # (B, 1, hidden_size)
+
+        state = None
+
+        for i in range(max_seq_length):
+            lstm_out, state = model.lstm.lstm(lstm_input, state)
+            output = model.lstm.fc(lstm_out.squeeze(1))
+
+            top_k_probs, top_k_indices = torch.topk(F.softmax(output, dim=1), top_k, dim=1)
+            top_k_probs = top_k_probs / torch.sum(top_k_probs, dim=1, keepdim=True)  
+            top_k_samples = torch.multinomial(top_k_probs, 1).squeeze()
+
+            predicted_word_indices = top_k_indices[range(batch_size), top_k_samples]
+
+            lstm_input = model.lstm.embedding(predicted_word_indices).unsqueeze(1)  # (B, 1, hidden_size)
+
+            for j in range(batch_size):
+                if end_token_appear[j]:
+                    continue
+
+                word = vocab.lookup_token(predicted_word_indices[j].item())
+                if word == end_token:
+                    end_token_appear[j] = True
+
+                captions[j].append(predicted_word_indices[j].item())
+
+        captions = [decoder(caption, vocab) for caption in captions]
+
+    return captions
+
+
+
+
+class ResNet50(nn.Module):
+    def __init__(self):
+        super(ResNet50, self).__init__()
+        self.ResNet50 = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+
+        self.ResNet50.fc = nn.Sequential(
+                            nn.Linear(2048, 1024),
+                            nn.ReLU(),
+                            nn.Dropout(0.5),
+                            nn.Linear(1024, 1024),
+                            nn.ReLU(),
+                           )
+
+        for k,v in self.ResNet50.named_parameters(recurse=True):
+          if 'fc' in k:
+            v.requires_grad = True
+          else:
+            v.requires_grad = False
+
+    def forward(self,x):
+        return self.ResNet50(x)       
+
+## lSTM (Decoder)
+
+class lstm(nn.Module):
+    def __init__(self, input_size, hidden_size, number_layers, embedding_dim, vocab_size):
+        super(lstm, self).__init__()
+
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.number_layers = number_layers
+        self.embedding_dim = embedding_dim
+        self.vocab_size = vocab_size
+        
+        self.embedding = nn.Embedding(vocab_size, hidden_size)
+        self.projection = nn.Linear(input_size, hidden_size)
+        self.relu = nn.ReLU()
+        
+        self.lstm = nn.LSTM(
+            input_size=hidden_size,
+            hidden_size=hidden_size,
+            num_layers=number_layers,
+            dropout=0.5,
+            batch_first=True,
+        )
+
+        self.fc = nn.Linear(hidden_size, vocab_size)
+
+    def forward(self, x, captions):
+        projected_image = self.projection(x).unsqueeze(dim=1)
+        embeddings = self.embedding(captions[:, :-1])
+    
+        # Concatenate the image feature as frist step with word embeddings
+        lstm_input = torch.cat((projected_image, embeddings), dim=1)
+        # print(torch.all(projected_image[:, 0, :] == lstm_input[:, 0, :])) # check
+        
+        lstm_out, _ = self.lstm(lstm_input)
+        logits = self.fc(lstm_out)
+
+        return logits
+
+## ImgCap
+
+class ImgCap(nn.Module):
+    def __init__(self, cnn_feature_size, lstm_hidden_size, num_layers, vocab_size, embedding_dim):
+        super(ImgCap, self).__init__()
+
+        self.cnn = ResNet50()
+
+        self.lstm = lstm(input_size=cnn_feature_size,
+                         hidden_size=lstm_hidden_size,
+                         number_layers=num_layers,
+                         embedding_dim=embedding_dim,
+                         vocab_size=vocab_size)
+
+    def forward(self, images, captions):
+        cnn_features = self.cnn(images)
+        output = self.lstm(cnn_features, captions)
+        return output
+
+    def generate_caption(self, images, vocab, decoder, device="cpu", start_token="<sos>", end_token="<eos>", max_seq_length=100):
+        self.eval()
+
+        with torch.no_grad():
+            start_index = vocab[start_token]
+            end_index = vocab[end_token]
+            images = images.to(device)
+            batch_size = images.size(0)
+
+            end_token_appear = {i: False for i in range(batch_size)}
+            captions = [[] for _ in range(batch_size)]
+
+            cnn_feature = self.cnn(images)
+            lstm_input = self.lstm.projection(cnn_feature).unsqueeze(1)  # (B, 1, hidden_size)
+
+            state = None
+
+            for i in range(max_seq_length):
+                lstm_out, state = self.lstm.lstm(lstm_input, state)
+                output = self.lstm.fc(lstm_out.squeeze(1))
+                predicted_word_indices = torch.argmax(output, dim=1)
+                lstm_input = self.lstm.embedding(predicted_word_indices).unsqueeze(1)  # (B, 1, hidden_size)
+
+                for j in range(batch_size):
+                    if end_token_appear[j]:
+                        continue
+
+                    word = vocab.lookup_token(predicted_word_indices[j].item())
+                    if word == end_token:
+                        end_token_appear[j] = True
+
+                    captions[j].append(predicted_word_indices[j].item())
+
+            captions = [decoder(caption) for caption in captions]
+
+        return captions
+