A custom image-captioning model trained from scratch on MS COCO, built around a fine-tuned CNN encoder and a Transformer decoder with cross-attention. The model generates natural-language descriptions of images, supports controllable caption length, produces interpretable attention heatmaps, and is benchmarked head-to-head against the pre-trained BLIP model.
Coursework for Deep Network Development (DNDEG), Eötvös Loránd University (ELTE), 2026 Spring.
- Encoder–Attention–Decoder architecture implemented in PyTorch from the ground up.
- Fine-tuned ResNet-18 encoder turning each image into a sequence of spatial tokens.
- Transformer decoder with masked self-attention and image cross-attention.
- Controllable generation via
<short>/<normal>/<long>control tokens. - Cross-attention visualization — heatmaps showing which image regions drive each word.
- BLEU-1…4 evaluation against all five reference captions per image.
- Baseline comparison with Salesforce BLIP.
┌────────────────────────┐
Image ─────▶│ CNN Encoder (ResNet-18)│──▶ spatial image tokens (B, 49, d_model)
│ fine-tuned, 1×1 proj │
└────────────────────────┘
│
▼ (cross-attention)
Caption ────▶┌────────────────────────┐
tokens │ Transformer Decoder │──▶ next-word distribution
+ control │ masked self-attention │
│ + cross-attention │
└────────────────────────┘
Encoder. A pre-trained ResNet-18 with its classification head removed. The final feature map is projected to the model dimension with a 1×1 convolution and flattened into a sequence of spatial tokens (7×7 = 49 tokens for 224×224 inputs). The backbone is fine-tuned during training.
Decoder. A Transformer decoder with sinusoidal positional encodings, masked self-attention over the caption tokens, and cross-attention over the encoder's image tokens. A control token prepended to the sequence steers the generated caption length.
MS COCO Captions 2017 — diverse everyday images, each with five human-written captions.
To fit within available compute, training uses a subset: 10,000 training images and
1,000 validation images. The vocabulary keeps words appearing at least 5 times, with
<pad>, <bos>, <eos>, <unk> special tokens (plus the length-control tokens). Captions are
capped at 30 tokens.
| Setting | Value |
|---|---|
| Loss | Cross-entropy (ignoring <pad>) |
| Optimizer | Adam, with a lower LR for the pre-trained encoder and a higher LR for the decoder |
| Regularization | Teacher forcing, gradient clipping (1.0), early stopping |
| Epochs | 5 (best checkpoint restored by validation loss) |
| Image size | 224 × 224, ImageNet normalization |
| Augmentation | Random resized crop, horizontal flip, color jitter (images only) |
Over 5 epochs the custom model's training loss fell from 4.12 → 2.70 and validation loss from 3.44 → 2.75, with no strong sign of overfitting.
BLEU scores (validation set, all 5 references):
| Metric | Custom model | Pre-trained BLIP |
|---|---|---|
| BLEU-1 | 0.645 | 0.648 |
| BLEU-2 | 0.461 | 0.524 |
| BLEU-3 | 0.315 | 0.406 |
| BLEU-4 | 0.211 | 0.307 |
The custom model is competitive on unigram overlap (BLEU-1) — it reliably names the right objects — but BLIP pulls ahead on longer n-grams, reflecting its large-scale vision-language pre-training and stronger fluency on multi-word phrases. Given that the custom model trains on a small COCO subset for only 5 epochs, matching BLIP on BLEU-1 is a solid outcome.
Qualitative outputs in the notebook include generated-vs-reference caption comparisons and cross-attention heatmaps.
image-captioning-coco/
├── image_captioning.ipynb # full pipeline: data, model, training, eval, comparison
├── requirements.txt
├── LICENSE
└── README.md
# 1. Clone
git clone https://github.com/TanhaR31/custom-image-captioning-model.git
cd custom-image-captioning-model
# 2. (Recommended) create a virtual environment
python -m venv .venv && source .venv/bin/activate # Windows: .venv\Scripts\activate
# 3. Install dependencies
pip install -r requirements.txtDownload the data (MS COCO 2017 — large):
# Images
wget http://images.cocodataset.org/zips/train2017.zip # ~18 GB
wget http://images.cocodataset.org/zips/val2017.zip # ~1 GB
# Annotations (captions)
wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip # ~241 MBExtract into the layout the notebook expects:
coco/
├── images/
│ ├── train2017/
│ └── val2017/
└── annotations/
├── captions_train2017.json
└── captions_val2017.json
Then set the DATA_ROOT path near the top of the notebook and run the cells top to bottom.
Compute note: a CUDA-capable GPU is strongly recommended for training. The notebook was developed with the dataset subset above; adjust
MAX_TRAIN_IMAGES/MAX_VAL_IMAGESto trade accuracy for speed.
Python · PyTorch & TorchVision · NLTK (BLEU) · Hugging Face Transformers (BLIP baseline) · Pillow · Matplotlib · NumPy
- Course: Deep Network Development, Faculty of Informatics, ELTE. Assignment authored by course staff.
- Dataset: MS COCO (Lin et al., 2014) — captions subset.
- Baseline: BLIP (Salesforce) via Hugging Face.
This repository contains my own solution, shared for learning and portfolio purposes. If you are currently taking this course, please follow your institution's academic-integrity policy and use it only as a reference.
Released under the MIT License.