Image Captioning

Modern image captioning stack for the COCO 2014 dataset. The repo bundles ViT/CNN encoders, RNN/LSTM/Transformer decoders, evaluation utilities, attention visualizations, and ready-to-run job scripts.

Environment Setup

1. Create a virtual environment and install dependencies:

   python3 -m venv .venv
   source .venv/bin/activate
   pip install -r requirements.txt

2. Download the COCO 2014 images and captions so the tree looks like:

   data/
     annotations/captions_{train,val}2014.json
     images/{train,val}2014/
   

1. Training(coco_cnn_rnn test → full run)

Quick run

python scripts/train_subset.py \
  --cocoapi-loc data \
  --tokenizer-name bert-base-uncased \
  --subset-size 64 \
  --epochs 1 \
  --max-steps 25 \
  --save-path artifacts/coco_cnn_rnn.pt \
  --log-dir runs/subset

Key flags:

• --encoder-arch {cnn,vit} and --decoder-arch {rnn,lstm,transformer} switch backbones/decoders.
• --train-backbone/--freeze-backbone control finetuning.
• --num-predictions and --eval-temperature collect quick metrics each epoch.

Cluster/production jobs

All SLURM-friendly launchers live in jobs/. Copy one (e.g. jobs/job_aspire.sh for ViT+LSTM or jobs/job_aspire_transformer_gpt2.sh for ViT+GPT2) and adapt queues, checkpoints, and hyperparameters before submitting.

2. Inference: captions & attention

Collage of validation captions

python scripts/infer_val.py \
  --checkpoint artifacts/coco_cnn_rnn.pt \
  --cocoapi-loc data \
  --num-samples 8 \
  --output-grid runs/val_preds/latest.png

The script reuses the tokenizer stored inside the checkpoint (override with --tokenizer-name). Pass --visualize-attention to dump per-token overlays next to the collage.

Single image caption

python scripts/caption_image.py \
  --checkpoint artifacts/coco_cnn_rnn.pt \
  --image path/to/image.jpg \
  --max-length 32 \
  --attention-dir runs/captions/attention

Gives a console caption and optional attention PNGs when --attention-dir is provided.

Visual attention maps for ViT+Transformer

python scripts/infer_attention_visualization.py \
  --checkpoint pretrain/coco_vit_transformer.pt \
  --cocoapi-loc data \
  --tokenizer-name gpt2 \
  --num-samples 4 \
  --output-dir runs/attention_viz

Outputs side-by-side overlays per token plus text summaries. Use --layer-to-viz to inspect a specific decoder layer and --max-viz-tokens / --tokens-per-row to control collage size.

3. Results

Attention Visualization

The model learns to attend to relevant image regions when generating each word:

Example: "a plane flying through the air with clouds in the background" - showing attention weights for each token

Validation Set Predictions

Sample predictions on COCO validation images:

Diverse captions generated for various scenes from the COCO dataset

4. TensorBoard logs

Training writes summaries under runs/. Point TensorBoard at the parent directory or a specific experiment:

tensorboard --logdir runs

The runs/ directory matches the value passed via --log-dir in scripts/train_subset.py.

5. Repository layout

• models/: encoders, decoders, and the CaptioningModel wrapper.
• dataset.py, dataloader.py: COCO glue + HuggingFace tokenization.
• scripts/: training, evaluation, inference, and visualization entry points.
• jobs/: cluster launchers relocated from the repo root.
• evaluation/: BLEU/CIDEr/METEOR metrics.
• artifacts/: checkpoints.