Image captioning with a ResNet18 encoder and a transformer decoder
An image captioning model built from scratch in PyTorch, trained on Flickr8k to a best validation perplexity of 19.74 — and a demo that runs the same idea backwards, retrieving images from a sentence.
The question
Cross-modal generation means conditioning on one modality and producing another: hand the model a photograph, and it returns a sentence. I wanted to build that end to end rather than call an API for it — to see what an encoder actually hands a decoder, and what the decoder does with it.
The task also has a natural mirror. If images and sentences can be placed in one shared space, the same relationship that lets you caption a picture should let you search pictures with a sentence. The project ended up covering both directions.
How it is put together
A ResNet18 pretrained on ImageNet is truncated before global pooling, so instead of collapsing the image to one vector it keeps a 7×7×512 feature map. A 1×1 convolution projects each cell to 256 dimensions, giving 49 "image tokens" that preserve where things are in the frame. The decoder cross-attends to those.
| Component | Detail |
|---|---|
| Image encoder | ResNet18, truncated before pooling → 7×7×512 → 1×1 conv → 49 tokens of 256-d |
| Tokenizer | BERT bert-base-uncased WordPiece, ~30k vocab; [CLS] as BOS, [SEP] as EOS |
| Decoder | 4-layer transformer decoder, norm_first, causal self-attention plus cross-attention |
| Output head | Linear to vocab logits, weight-tied to the input token embedding |
| Loss | Cross-entropy on next-token prediction, padding ignored |
Training uses teacher forcing — the decoder sees ground-truth tokens shifted right and predicts the next one at every position. At inference it generates greedily, one token at a time, stopping at [SEP] or a length cap.
What the training showed
The dataset is Flickr8k: 8,091 images with five human-written captions each, so roughly 40,000 image-caption pairs. Training ran for 20 epochs at about 53 seconds per epoch on an NVIDIA H100.
Best validation loss was 2.9827 at epoch 12, a perplexity of 19.74. Perplexity is the more interpretable number: on average the model is choosing the correct next word from about twenty plausible candidates out of a 30,000-word vocabulary. From epoch 13 onward validation loss creeps back up while training loss keeps falling — textbook mild overfitting, and the reason the epoch-12 checkpoint is the one kept.
Running it backwards
The demo pairs the captioning model with CLIP ViT-B/32, a pretrained joint image-text encoder. Image embeddings are computed once and cached, then a text query is embedded into the same space and images are ranked by cosine similarity. The captioning model describes the top matches.
So one interface does both directions: give it a picture and it writes a sentence, or give it a sentence and it finds pictures and then describes what it found.