A complete multimodal project — building an image captioning model that combines a CNN vision encoder with an RNN text decoder, a hands-on introduction to connecting two different modalities within a single trainable architecture.
Problem Statement
Build a model that generates a natural-language caption describing the content of an input image, combining computer vision and sequence generation into a single end-to-end system.
Dataset
A dataset of images paired with human-written captions — Flickr8k (8,000 images, 5 captions each) is a well-suited, manageable size for a first captioning project.
Architecture & Approach
A pretrained CNN (e.g. ResNet, with its final classification layer removed) encodes the image into a feature vector; an LSTM decoder then generates the caption word by word, conditioned on the image feature vector — directly combining the CNN Fundamentals and RNN/LSTM categories into one working system.
Step-by-Step Build
import torch
import torch.nn as nn
import torchvision.models as models
class ImageEncoder(nn.Module):
def __init__(self, embed_dim=256):
super().__init__()
resnet = models.resnet50(weights='IMAGENET1K_V2')
self.backbone = nn.Sequential(*list(resnet.children())[:-1]) # remove the final classification layer
for param in self.backbone.parameters():
param.requires_grad = False # freeze the pretrained CNN backbone
self.projection = nn.Linear(resnet.fc.in_features, embed_dim)
def forward(self, images):
features = self.backbone(images).squeeze(-1).squeeze(-1) # (batch, 2048)
return self.projection(features) # (batch, embed_dim)
class CaptionDecoder(nn.Module):
def __init__(self, vocab_size, embed_dim=256, hidden_dim=512):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
self.fc = nn.Linear(hidden_dim, vocab_size)
def forward(self, image_features, captions):
embedded_captions = self.embedding(captions)
# Use the image feature vector as the LSTM's initial input, before the caption words
inputs = torch.cat([image_features.unsqueeze(1), embedded_captions], dim=1)
lstm_out, _ = self.lstm(inputs)
return self.fc(lstm_out)
encoder = ImageEncoder()
decoder = CaptionDecoder(vocab_size=len(vocab))
optimizer = torch.optim.Adam(list(encoder.projection.parameters()) + list(decoder.parameters()), lr=0.0003)
loss_fn = nn.CrossEntropyLoss(ignore_index=vocab["<pad>"])
for epoch in range(15):
for images, captions in train_loader:
image_features = encoder(images)
outputs = decoder(image_features, captions[:, :-1]) # feed all but the last token as input
loss = loss_fn(outputs.reshape(-1, len(vocab)), captions[:, :].reshape(-1))
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(f"Epoch {epoch+1}: loss={loss.item():.4f}")
# Generate a caption for a new image, autoregressively
@torch.no_grad()
def generate_caption(image, max_len=20):
image_features = encoder(image.unsqueeze(0))
caption = [vocab["<start>"]]
hidden = None
inputs = image_features.unsqueeze(1)
for _ in range(max_len):
lstm_out, hidden = decoder.lstm(inputs, hidden)
logits = decoder.fc(lstm_out.squeeze(1))
next_word = logits.argmax(dim=-1).item()
if next_word == vocab["<end>"]: break
caption.append(next_word)
inputs = decoder.embedding(torch.tensor([[next_word]]))
return ' '.join(idx_to_word[idx] for idx in caption[1:])Expected Results
After sufficient training, expect captions that correctly identify major objects and simple actions in most test images (e.g. "a dog running on grass"), though fine-grained detail and unusual scenes will often be handled less reliably than a large, modern captioning model — a fair and expected result for a model of this modest scale and training budget.
Key Learnings & Extensions
- This project makes the "connecting a pretrained vision encoder to a text-generating decoder" pattern from Vision-Language Models concrete at a small, tractable scale — the same fundamental idea (though at vastly larger scale, with a Transformer decoder) powers modern VLMs.
- Extension: Add an attention mechanism so the decoder attends to different spatial regions of the image feature map at each generation step (rather than a single fixed feature vector) — this is the "Show, Attend and Tell" style approach, and produces noticeably better, more grounded captions.
- Extension: Evaluate captions quantitatively using BLEU or ROUGE score against the reference captions, applying the metrics from BLEU Score.
