model.py

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models import resnet50, ResNet50_Weights

class EncoderCNN(nn.Module):
    def __init__(self, embed_size):
        """
        Load the pretrained ResNet-50 and replace top fc layer.
        
        Args:
            embed_size: dimensionality of image embeddings
        """
        super(EncoderCNN, self).__init__()
        
        # Load pretrained ResNet-50
        resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
        
        # Freeze all ResNet parameters (we don't want to update them)
        for param in resnet.parameters():
            param.requires_grad_(False)
        
        # Remove the last fully-connected layer
        modules = list(resnet.children())[:-1]
        self.resnet = nn.Sequential(*modules)
        
        # Replace with a new embedding layer
        self.embed = nn.Linear(resnet.fc.in_features, embed_size)
    
    def forward(self, images):
        """
        Extract feature vectors from input images.
        
        Args:
            images: batch of images (batch_size, 3, 224, 224)
        
        Returns:
            features: image embeddings (batch_size, embed_size)
        """
        # Pass images through ResNet
        features = self.resnet(images)
        
        # Flatten the features
        features = features.view(features.size(0), -1)
        
        # Pass through embedding layer
        features = self.embed(features)
        
        return features


class DecoderRNN(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers=1):
        """
        Initialize the DecoderRNN.
        
        Args:
            embed_size: dimensionality of image and word embeddings
            hidden_size: number of features in hidden state of the RNN
            vocab_size: size of vocabulary
            num_layers: number of layers in LSTM
        """
        super(DecoderRNN, self).__init__()
        
        self.embed_size = embed_size
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size
        self.num_layers = num_layers
        
        # Embedding layer: converts word indices to embeddings
        self.embed = nn.Embedding(vocab_size, embed_size)
        
        # LSTM layer: processes sequences
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
        
        # Linear layer: maps LSTM output to vocabulary scores
        self.fc = nn.Linear(hidden_size, vocab_size)
    
    def forward(self, features, captions):
        """
        Forward pass during training.
        
        Args:
            features: image features from encoder (batch_size, embed_size)
            captions: ground-truth captions (batch_size, caption_length)
        
        Returns:
            outputs: predicted word scores (batch_size, caption_length, vocab_size)
        """
        # Remove the last token (<end>) from captions for input
        # We predict the next word, so we don't need <end> as input
        captions = captions[:, :-1]
        
        # Convert caption word indices to embeddings
        embeddings = self.embed(captions)  # (batch_size, caption_length-1, embed_size)
        
        # Concatenate image features as the first time step
        # features.unsqueeze(1) adds a time dimension: (batch_size, 1, embed_size)
        embeddings = torch.cat((features.unsqueeze(1), embeddings), dim=1)
        
        # Pass through LSTM
        lstm_out, _ = self.lstm(embeddings)  # (batch_size, caption_length, hidden_size)
        
        # Map to vocabulary scores
        outputs = self.fc(lstm_out)  # (batch_size, caption_length, vocab_size)
        
        return outputs
    
    def sample(self, inputs, states=None, max_len=20):
        """
        Accepts pre-processed image tensor (inputs) and returns predicted 
        sentence (list of tensor ids of length max_len).
        
        Args:
            inputs: image features (1, embed_size)
            states: initial hidden state (optional)
            max_len: maximum length of generated caption
        
        Returns:
            predicted_ids: list of predicted word indices
        """
        predicted_ids = []
        
        # inputs should have shape (1, 1, embed_size) for LSTM
        inputs = inputs.unsqueeze(1)
        
        for i in range(max_len):
            # Pass through LSTM
            lstm_out, states = self.lstm(inputs, states)  # lstm_out: (1, 1, hidden_size)
            
            # Pass through linear layer
            outputs = self.fc(lstm_out.squeeze(1))  # outputs: (1, vocab_size)
            
            # Get the word with highest score
            _, predicted = outputs.max(1)  # predicted: (1)
            
            # Store the predicted word index
            predicted_ids.append(predicted.item())
            
            # Use the predicted word as input for next time step
            inputs = self.embed(predicted)  # (1, embed_size)
            inputs = inputs.unsqueeze(1)    # (1, 1, embed_size)
        
        return predicted_ids
class DecoderRNN(nn.Module):
    def __init__(self, embed_size, hidden_size, vocab_size, num_layers=1):
        """Initialize the DecoderRNN."""
        super(DecoderRNN, self).__init__()  # CRITICAL: Initialize parent class
        
        self.embed_size = embed_size
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size
        self.num_layers = num_layers
        
        # Embedding layer: converts word indices to embeddings
        self.embed = nn.Embedding(vocab_size, embed_size)
        
        # LSTM layer: processes sequences
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
        
        # Linear layer: maps LSTM output to vocabulary scores
        self.fc = nn.Linear(hidden_size, vocab_size)
    
    def forward(self, features, captions):
        """
        Forward pass during training.
        
        Args:
            features: image features from encoder (batch_size, embed_size)
            captions: ground-truth captions (batch_size, caption_length)
        
        Returns:
            outputs: predicted word scores (batch_size, caption_length, vocab_size)
        """
        # Remove the last token (<end>) from captions for input
        # We predict the next word, so we don't need <end> as input
        captions = captions[:, :-1]
        
        # Convert caption word indices to embeddings
        embeddings = self.embed(captions)  # (batch_size, caption_length-1, embed_size)
        
        # Concatenate image features as the first time step
        # features.unsqueeze(1) adds a time dimension: (batch_size, 1, embed_size)
        embeddings = torch.cat((features.unsqueeze(1), embeddings), dim=1)
        
        # Pass through LSTM
        lstm_out, _ = self.lstm(embeddings)  # (batch_size, caption_length, hidden_size)
        
        # Map to vocabulary scores
        outputs = self.fc(lstm_out)  # (batch_size, caption_length, vocab_size)
        
        return outputs
    
    def sample(self, inputs, states=None, max_len=20):
        """
        Accepts pre-processed image tensor (inputs) and returns predicted 
        sentence (list of tensor ids of length max_len).
        
        Args:
            inputs: image features (1, embed_size)
            states: initial hidden state (optional)
            max_len: maximum length of generated caption
        
        Returns:
            predicted_ids: list of predicted word indices
        """
        predicted_ids = []
        
        # inputs should have shape (1, 1, embed_size) for LSTM
        inputs = inputs.unsqueeze(1)
        
        for i in range(max_len):
            # Pass through LSTM
            lstm_out, states = self.lstm(inputs, states)  # lstm_out: (1, 1, hidden_size)
            
            # Pass through linear layer
            outputs = self.fc(lstm_out.squeeze(1))  # outputs: (1, vocab_size)
            
            # Get the word with highest score
            _, predicted = outputs.max(1)  # predicted: (1)
            
            # Store the predicted word index
            predicted_ids.append(predicted.item())
            
            # Use the predicted word as input for next time step
            inputs = self.embed(predicted)  # (1, embed_size)
            inputs = inputs.unsqueeze(1)    # (1, 1, embed_size)
        
        return predicted_ids