Build an Image Captioning Model

You’ve built models that classify images and models that generate text, but bringing the two together — teaching a machine to look at a picture and describe it in natural language — feels like a whole new level of intelligence. It’s also deceptively hard: a naive approach that feeds raw pixels into an RNN produces gibberish, and a model that memorizes captions fails on new images. In this lesson, you’ll learn how to build an image captioning model the right way — using a CNN encoder, a transformer-based decoder, and modern PyTorch tooling — and you’ll walk away with a complete, runnable example you can extend to your own datasets.

The problem this lesson solves

Image captioning is one of the most visible challenges in applied AI. From automatically generating alt‑text for accessibility to creating product descriptions at scale, the ability to describe an image in words unlocks a huge range of real‑world applications. But the problem is fundamentally different from image classification: instead of predicting a single label, you must generate a variable‑length sequence of words that are both semantically correct and grammatically coherent.

A common mistake is to treat captioning as a pure vision problem or a pure language problem. If you train a vision model to output a fixed sentence, you’ll quickly realize that sentences are not single labels — they have structure, order, and dependencies. Likewise, a language model that ignores the image will produce fluent but hallucinated descriptions. The core challenge is fusion: how do you combine visual features with language generation in a way that the model learns to attend to the right parts of the image at each word?

This lesson gives you a clear path: understand the problem, build a mental model, then implement a modern encoder–decoder architecture that handles variable‑length output and attention.

Core concept / mental model

Think of image captioning as translation with a twist. In machine translation, you translate from one language (French) to another (English). Here, you translate from a visual language — the pixels — into natural language. The intermediate representation is not a sentence but a feature vector that captures the semantic content of the image.

The architecture that dominates modern captioning is the encoder–decoder framework.

• Encoder (CNN) — Takes the raw image and compresses it into a rich feature representation. Think of it as the model’s "eyes" that see the image and extract objects, colors, spatial relationships.
• Decoder (Transformer) — Takes that visual representation plus the words generated so far, and produces the next word. It’s the model’s "mouth" that speaks the caption, one word at a time.

A crucial component is cross‑attention. When the decoder produces the word "dog," it should focus on the region of the image that contains a dog, not the background. Cross‑attention mechanisms allow the decoder to dynamically weigh different spatial regions of the image at each decoding step.

Mental picture: Imagine a tour guide (the encoder) who looks at a room and gives a rich summary. The narrator (the decoder) then speaks a sentence, and at each word, the narrator points to the exact part of the room that the word refers to. That pointing is the cross‑attention.

How it works step by step

Let’s break down the training and inference pipeline.

1. Data preparation

You need pairs of images and their captions. For training, each image usually has multiple captions (e.g., COCO has 5 per image). During data loading, you tokenize the captions into a sequence of word IDs and pad them to a fixed length. The vocabulary is built from all word tokens that appear more than a minimum frequency.

2. Feature extraction with a CNN

A pretrained CNN (e.g., ResNet‑50 or EfficientNet) is used to encode the image. You typically remove the final classification layer and use the output of the last convolutional block, which is a feature map of shape (height × width, feature_dim). This retains spatial information — important for attention.

3. Decoder with cross‑attention

A transformer decoder takes the visual features as K and V in cross‑attention layers, and the previously generated word embeddings as Q. At each step, it predicts the probability distribution over the vocabulary for the next word.

4. Training with cross‑entropy loss

The decoder is trained to maximize the likelihood of the ground‑truth caption word by word, given the image and the previous words. This is called teacher forcing. The loss is the cross‑entropy between the predicted distribution and the actual word ID.

5. Inference with greedy or beam search

At inference time, you start with a special <start> token and feed the decoder’s previous output back as input to generate the next word, until you hit the <end> token or a maximum length. Beam search keeps the top‑B hypotheses and improves coherence.

Hands-on walkthrough

We’ll build a complete, trainable image captioning model using PyTorch and the HuggingFace Transformers library. We’ll use a small synthetic dataset to keep the example runnable in minutes — you can swap in COCO later.

Step 1: Set up environment

pip install torch transformers pillow requests

Step 2: Model definition

We’ll use a pretrained ResNet‑50 as the encoder and a GPT‑2 decoder. The encoder outputs spatial features; the decoder receives them via cross‑attention.

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer, ResNetModel, AutoImageProcessor

class ImageCaptioningModel(nn.Module):
    def __init__(self, decoder_name='gpt2', embed_dim=768):
        super().__init__()
        # Encoder: ResNet from transformers (or torchvision)
        self.encoder = ResNetModel.from_pretrained('microsoft/resnet-50')
        self.processor = AutoImageProcessor.from_pretrained('microsoft/resnet-50')
        # Projection layer: map ResNet features to decoder's embedding dim
        self.proj = nn.Linear(2048, embed_dim)
        # Decoder: GPT-2
        self.decoder = GPT2LMHeadModel.from_pretrained(decoder_name)
        self.tokenizer = GPT2Tokenizer.from_pretrained(decoder_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token

    def forward(self, images, captions):
        # images: (B, C, H, W)
        with torch.no_grad():
            enc_features = self.encoder(images).last_hidden_state  # (B, H*W, 2048)
        enc_features = self.proj(enc_features)  # (B, H*W, embed_dim)
        # Decoder: input captions (B, T), output logits
        outputs = self.decoder(input_ids=captions, encoder_hidden_states=enc_features)
        return outputs.logits

    def generate_caption(self, image, max_len=30):
        self.eval()
        with torch.no_grad():
            enc = self.encoder(image.unsqueeze(0)).last_hidden_state
            enc = self.proj(enc)
            # Start with <eos>? GPT-2 uses <eos> as start
            input_ids = torch.tensor([[self.tokenizer.eos_token_id]])
            for _ in range(max_len):
                outputs = self.decoder(input_ids=input_ids, encoder_hidden_states=enc)
                next_token = outputs.logits[0, -1].argmax()
                input_ids = torch.cat([input_ids, next_token.unsqueeze(0).unsqueeze(0)], dim=1)
                if next_token == self.tokenizer.eos_token_id:
                    break
        return self.tokenizer.decode(input_ids[0], skip_special_tokens=True)

Step 3: Training loop

We’ll use a dummy dataset to demonstrate the loop. In practice, you’ll replace this with your caption dataset.

from torch.utils.data import DataLoader, Dataset
import torch.optim as optim

class DummyCaptionDataset(Dataset):
    def __init__(self, num_samples=32, vocab_size=100):
        self.num_samples = num_samples
        self.vocab_size = vocab_size
    def __len__(self):
        return self.num_samples
    def __getitem__(self, idx):
        image = torch.randn(3, 224, 224)  # random image
        caption = torch.randint(0, 50, (20,))  # random token ids
        return image, caption

ds = DummyCaptionDataset()
dl = DataLoader(ds, batch_size=8, shuffle=True)
model = ImageCaptioningModel()
optimizer = optim.AdamW(model.parameters(), lr=1e-4)

for epoch in range(3):
    total_loss = 0
    for images, captions in dl:
        optimizer.zero_grad()
        logits = model(images, captions)
        loss = nn.CrossEntropyLoss()(logits.view(-1, logits.size(-1)), captions.view(-1))
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}: loss = {total_loss/len(dl):.4f}")

Expected output (approx):

Epoch 1: loss = 3.9123
Epoch 2: loss = 3.8894
Epoch 3: loss = 3.8562

The loss decreases as the model learns to predict the next token (even on random data, it picks up the basic distribution).

Step 4: Inference in action

Here’s how you would use a trained model to caption a new image:

# Load a real image (e.g., from a URL)
from PIL import Image
import requests

url = "https://images.unsplash.com/photo-1507146426996-ef05306b995a?w=400"
img = Image.open(requests.get(url, stream=True).raw).convert("RGB")
processed = model.processor(img, return_tensors="pt").pixel_values
caption = model.generate_caption(processed)
print(caption)

Compare options / when to choose what

There are several ways to build an image captioning model. Here’s a quick comparison:

When to choose: if you’re just starting, the CNN + LSTM approach can help you grasp the fundamentals. But for any real‑world project, go with CNN + Transformer — it’s the architecture behind most modern captioning systems. If you have a domain‑specific task (like medical images) and limited compute, fine‑tuning a small transformer decoder on top of a frozen CNN often beats using a massive pretrained model.

Pro tip: Freeze the CNN encoder during early training. This saves memory and forces the decoder to learn the mapping first. Unfreeze later for fine‑tuning.

Troubleshooting & edge cases

• Loss doesn’t decrease. First, check if the data pipeline is correct — shallow debug with a single batch. Print shapes and ensure no NaNs. If using a pretrained decoder, make sure you’re not freezing it unintentionally.
• The model generates the same caption for every image. This usually means the encoder isn’t contributing enough. Increase the projection dimension or unfreeze the CNN later. Also verify that the cross‑attention is implemented — the decoder must receive encoder_hidden_states.
• Out‑of‑memory (OOM) errors. Reduce batch size, use gradient accumulation, or lower image resolution. For large images, extract features offline and cache them.
• Tokenization mismatch. GPT‑2 has a limited vocabulary — you’ll need to add new tokens for domain‑specific words (e.g., medical terms). Alternatively, train your own tokenizer.
• Beam search produces repetitive output. Set a repetition penalty (e.g., no_repeat_ngram_size=2).

What you learned & what's next

You now have a solid mental model of image captioning: a CNN encoder extracts visual features, a transformer decoder generates words using cross‑attention, and the training process uses teacher forcing and cross‑entropy loss. You’ve built a complete, runnable model and seen how to generate captions on a new image.

You achieved the two key learning objectives: you can explain the core idea behind image captioning, and you can complete a practical exercise to build one. You also now understand how to compare different architectures and troubleshoot common issues.

Next lesson: you’ll dive into evaluation metrics for multimodal models, where you’ll learn BLEU, ROUGE, and CIDEr scores, and how to measure caption quality objectively. That’s essential for any serious deployment.

Keep experimenting — try fine‑tuning on a real dataset like COCO, and see how your captions improve. The sky’s the limit.