Spaces:
Running
Running
| import torch | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| from typing import Tuple | |
| class Encoder(nn.Module): | |
| """ | |
| 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 | |
| # 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 the last layer of ResNet-50 | |
| modules = list(resnet.children())[:-1] | |
| self.resnet = nn.Sequential(*modules) | |
| # 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 pass through the encoder. | |
| Args: | |
| images (torch.Tensor): Input images of shape (BATCH, 3, 224, 224). | |
| Returns: | |
| torch.Tensor: Image features of shape (BATCH, IMAGE_EMB_DIM). | |
| """ | |
| features = self.resnet(images) | |
| # Reshape features to (BATCH, 2048) | |
| features = features.reshape(features.size(0), -1).to(self.device) | |
| # Pass features through final linear layer | |
| features = self.fc(features).to(self.device) | |
| 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, | |
| word_emb_dim: int, | |
| hidden_dim: int, | |
| num_layers: int, | |
| vocab_size: int, | |
| device: torch.device): | |
| super(Decoder, self).__init__() | |
| self.word_emb_dim = word_emb_dim | |
| self.hidden_dim = hidden_dim | |
| self.num_layers = num_layers | |
| self.vocab_size = vocab_size | |
| self.device = device | |
| # 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))) | |
| # 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(self.hidden_dim, self.vocab_size), | |
| nn.LogSoftmax(dim=2) | |
| ) | |
| def forward(self, | |
| embedded_captions: torch.Tensor, | |
| hidden: torch.Tensor, | |
| cell: torch.Tensor) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: | |
| """ | |
| Forward pass through the decoder. | |
| Args: | |
| 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: | |
| 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. | |
| """ | |
| # Pass through LSTM | |
| output, (hidden, cell) = self.lstm(embedded_captions, (hidden, cell)) | |
| # Pass through final linear layer | |
| output = self.fc(output) | |
| return output, (hidden, cell) | |