Image Captioning Model
This repository contains a custom PyTorch image captioning model. The model receives an input image and generates a natural-language caption describing the image.
The architecture is built from two main components:
Image Encoder : EfficientNet-V2-S backbone pretrained on ImageNet.Text Decoder : Transformer decoder that generates captions token by token.
The model was trained for image caption generation using COCO-style image-caption pairs.
Model Architecture
The model follows an encoder-decoder structure:
1 Input Image
2 ↓
3 EfficientNet-V2-S Image Encoder
4 ↓
5 Image Feature Tokens
6 ↓
7 Transformer Text Decoder
8 ↓
9 Generated Caption
Image Encoder
The encoder uses EfficientNet_V2_S from torchvision.models.
The image encoder extracts visual features from the input image and projects them into a 256-dimensional embedding space. The final image representation is treated as a sequence of visual tokens.
Encoder details:
1 Backbone: EfficientNet-V2-S
2 Input image size: 224 x 224
3 Output visual tokens: 49
4 Embedding dimension: 256
5 ImageNet normalization: Yes
Text Decoder
The decoder is a Transformer decoder that generates captions autoregressively.
Decoder details:
1 Vocabulary size: 9,721
2 Embedding dimension: 256
3 Number of Transformer decoder layers: 6
4 Number of attention heads: 8
5 Feed-forward dimension: 1024
6 Maximum caption length: 52
7 Dropout: 0.1
8 Decoding methods: Greedy search and beam search
Repository Files
This repository contains:
1 best_phase1.pt # PyTorch checkpoint
2 Training-5k.ipynb # Training and inference notebook
The checkpoint contains:
Checkpoint information:
1 Checkpoint file: best_phase1.pt
2 Epoch: 8
3 Validation loss: 3.6158
Important Note About Vocabulary
This model uses a custom word-level vocabulary built from the training captions.
The checkpoint stores the model weights, but it does not store the vocabulary mapping. To run inference correctly, you must use the same vocabulary that was used during training.
The vocabulary contains 9,721 tokens and uses the following special tokens:
1 <PAD> = 0
2 <SOS> = 1
3 <EOS> = 2
4 <UNK> = 3
If you want to make this model easier to use, it is recommended to upload an additional file such as:
containing the stoi and itos mappings.
Training Data
The model was trained using COCO-style image-caption data.
The training notebook is configured to use:
1 Dataset format: COCO captions
2 Training annotations: captions_train2014.json
3 Validation annotations: captions_val2014.json
4 Image size: 224 x 224
5 Batch size: 32
6 Maximum caption length: 52
The notebook version included in this repository was designed for a smaller training experiment using a limited number of samples.
Image Preprocessing
Images are resized to 224 x 224 and normalized with ImageNet statistics:
1 IMAGENET_MEAN = [ 0.485 , 0.456 , 0.406 ]
2 IMAGENET_STD = [ 0.229 , 0.224 , 0.225 ]
Validation and inference transforms:
1 import torchvision . transforms as T
2
3 transform = T . Compose ( [
4 T . Resize ( ( 224 , 224 ) ) ,
5 T . ToTensor ( ) ,
6 T . Normalize (
7 mean = [ 0.485 , 0.456 , 0.406 ] ,
8 std = [ 0.229 , 0.224 , 0.225 ]
9 ) ,
10 ] )
How to Use
This is a custom PyTorch model. It is not a standard Hugging Face Transformers model, so it cannot be loaded directly with:
AutoModel.from_pretrained(...)To use the model, open and run the notebook:
The notebook contains:
1 Vocabulary class
2 Dataset class
3 EfficientNet encoder
4 Transformer decoder
5 ImageCaptioningModel class
6 Training loop
7 Checkpoint loading
8 Greedy decoding
9 Beam-search decoding
10 Evaluation code
Loading the Checkpoint
After defining the model architecture and rebuilding/loading the same vocabulary, the checkpoint can be loaded as follows:
1 import torch
2
3 device = torch . device ( "cuda" if torch . cuda . is_available ( ) else "cpu" )
4
5 model = ImageCaptioningModel (
6 vocab_size = 9721 ,
7 embed_dim = 256 ,
8 num_heads = 8 ,
9 num_layers = 6 ,
10 ff_dim = 1024 ,
11 max_len = 52 ,
12 dropout = 0.1
13 ) . to ( device )
14
15 checkpoint = torch . load ( "best_phase1.pt" , map_location = device )
16 model . load_state_dict ( checkpoint [ "model" ] )
17 model . eval ( )
18
19 print ( "Loaded checkpoint" )
20 print ( "Epoch:" , checkpoint [ "epoch" ] )
21 print ( "Validation loss:" , checkpoint [ "val_loss" ] )
Generating a Caption
The notebook includes two caption generation methods:
1 model . generate_greedy ( image_tensor )
2 model . generate_beam ( image_tensor , beam_size = 5 )
Example:
1 from PIL import Image
2
3 image = Image . open ( "example.jpg" ) . convert ( "RGB" )
4 image_tensor = transform ( image )
5
6 caption = model . generate_beam ( image_tensor , beam_size = 5 )
7 print ( caption )
Example Output
Example caption format:
a bicycle with a clock as the front wheelActual output quality depends on the training data size, checkpoint version, and decoding method.
Evaluation
The notebook includes BLEU evaluation code using NLTK:
from nltk.translate.bleu_score import corpus_bleu, SmoothingFunctionYou can evaluate the model on validation images using greedy decoding or beam search.
Recommended metrics for this task:
1 BLEU-1
2 BLEU-4
3 CIDEr
4 METEOR
5 ROUGE-L
Limitations
This model is an experimental image captioning model.
Known limitations:
The model uses a custom word-level tokenizer, not a subword tokenizer.
The vocabulary must match the original training vocabulary.
The checkpoint alone is not enough for fully reproducible inference unless the vocabulary is also available.
Caption quality may be limited if the model was trained on a small subset of the dataset.
The model may generate generic or repetitive captions.
The model may fail on images that are very different from the training distribution.
The model may hallucinate objects that are not present in the image.
Recommended Improvements
To make this repository easier to use, future versions should include:
1 vocab.json
2 model.py
3 requirements.txt
4 inference.py
5 example images
6 evaluation results
A better repository structure would be:
1 .
2 ├── README.md
3 ├── best_phase1.pt
4 ├── Training-5k.ipynb
5 ├── vocab.json
6 ├── model.py
7 ├── inference.py
8 └── requirements.txt
Requirements
The notebook uses the following main libraries:
1 torch
2 torchvision
3 Pillow
4 numpy
5 matplotlib
6 nltk
7 pycocotools
8 pycocoevalcap
9 einops
Install dependencies with:
pip install torch torchvision pillow numpy matplotlib nltk pycocotools pycocoevalcap einopsCitation
If you use this model, please cite or mention this repository.
Author
Created as a custom PyTorch image captioning model using an EfficientNet image encoder and a Transformer text decoder.
