anmol-unitmole/image-caption-generation-vision-encoder-decoder-model

Fine-Tuned ViT-GPT-2 Vision Encoder-Decoder Image Captioning

Model repository

anmol-unitmole/image-caption-generation-vision-encoder-decoder-model

System

image-caption-generation-vision-encoder-decoder-model

Status

Training, held-out evaluation, browser-compatible ONNX export, dynamic quantization, and PyTorch-versus-ONNX parity validation are complete.

This repository contains the browser-compatible model configuration, processor, tokenizer, generation configuration, and quantized ONNX artifacts for a fine-tuned ViT-GPT-2 Vision Encoder-Decoder image-captioning model.

The model was fine-tuned locally using an NVIDIA GeForce RTX 5090. The browser artifacts are intended for Transformers.js and ONNX Runtime Web inference inside a static Hugging Face Space.

The final public Static Space remains subject to successful browser loading and end-to-end caption-generation validation.

Architecture

• Base checkpoint: nlpconnect/vit-gpt2-image-captioning
• Vision encoder: Vision Transformer — ViT
• Text decoder: GPT-2 with encoder cross-attention
• Python framework: PyTorch
• Hugging Face architecture: VisionEncoderDecoderModel
• Pretrained browser baseline: Xenova/vit-gpt2-image-captioning
• Fine-tuned browser repository: anmol-unitmole/image-caption-generation-vision-encoder-decoder-model
• Browser runtime: Transformers.js and ONNX Runtime Web
• Preferred browser provider: WebGPU
• Fallback browser provider: WebAssembly

Task

The model generates English-language captions from input images.

Supported project features include:

• greedy caption generation;
• beam-search caption generation;
• caption-length controls;
• alternate-text draft generation;
• pretrained-versus-fine-tuned comparison;
• caption-structure diagnostics;
• genuine precomputed token-generation traces for bundled sample images;
• model latency and evaluation reporting.

Training method

The model was trained using a staged adaptation strategy.

Stage 1 — Decoder warm-up

• Freeze the complete vision encoder.
• Train the decoder and cross-modal generation pathway.
• Preserve pretrained visual representations during initial adaptation.

Stage 2 — Upper encoder adaptation

• Unfreeze the upper eight ViT encoder blocks.
• Use a lower learning rate for the vision encoder.
• Continue adapting the decoder using a higher learning rate.
• Evaluate generated captions during training.

Stage 3 — Conservative full fine-tuning

• Unfreeze the full encoder-decoder architecture.
• Apply reduced learning rates.
• Select the best checkpoint using validation CIDEr with configured fallback behavior.

Training controls included:

• BF16 mixed precision;
• TF32 support;
• fused AdamW fallback;
• gradient accumulation;
• label smoothing;
• warm-up scheduling;
• cosine learning-rate decay;
• minimum learning-rate floor;
• gradient clipping;
• early-stopping support;
• resumable checkpoints;
• best-checkpoint selection;
• environment and GPU-memory reporting.

Hardware

The final experiment was executed locally using:

• GPU: NVIDIA GeForce RTX 5090
• GPU memory: approximately 31.84 GB
• CUDA build: CUDA 13.0
• PyTorch: 2.12.0 with CUDA 13.0
• Precision: BF16 where supported
• Operating system: Windows 11

Dataset

The primary experiment used Flickr30k.

The processed experiment contained:

Split	Unique images	Caption records
Training	29,000	144,767
Validation	1,014	5,062
Test	1,000	4,985

Multiple reference captions remained grouped for evaluation.

The complete Flickr30k dataset is not redistributed through this model repository. Users must obtain the dataset according to its original access and licensing requirements.

Inputs

Supported browser image formats:

• JPG/JPEG
• PNG
• WEBP
• BMP

Input processing includes:

• safe image decoding;
• EXIF orientation correction;
• RGB conversion;
• model-native image resizing and normalization;
• processor-compatible pixel-value generation;
• validation for unsupported or corrupt images.

Outputs

The model may provide:

• one greedy caption;
• multiple ranked beam-search candidates;
• a concise alternate-text draft;
• generation latency;
• generation configuration;
• sequence score where available;
• genuine token-transition traces for bundled sample images;
• caption-structure diagnostics.

Generated captions are model predictions and are not guaranteed to be factually correct.

Evaluation protocol

The pretrained baseline and fine-tuned model were evaluated on the same held-out 1,000-image Flickr30k test subset.

Main results

Metric	Pretrained baseline	Fine-tuned model
BLEU-1	0.5779	0.4350
BLEU-4	0.1366	0.1078
METEOR	0.3337	0.4050
ROUGE-L	0.3759	0.3041
CIDEr	0.2336	0.0595
CLIPScore	0.6601	0.7211
Distinct-1	0.0652	0.0708
Distinct-2	0.2331	0.2478
Unique-caption ratio	0.9500	1.0000
Average repetition rate	0.1126	0.1463
Mean GPU latency	0.1163 seconds	0.2762 seconds

Result interpretation

The fine-tuned model did not outperform the pretrained baseline on every metric.

Observed improvements included:

• higher METEOR;
• higher CLIPScore;
• higher Distinct-1;
• higher Distinct-2;
• a 1.0 unique-caption ratio.

Observed regressions included:

• lower BLEU-1;
• lower BLEU-4;
• lower ROUGE-L;
• lower CIDEr;
• increased caption length;
• increased repetition;
• slower generation latency.

These mixed results are reported intentionally. The experiment demonstrates that fine-tuning can improve semantic similarity and lexical diversity while reducing reference-overlap metrics or generation efficiency.

The fine-tuned model should therefore not be described as universally superior to the pretrained baseline.

Model size

The fine-tuned Python checkpoint contains approximately:

• Parameters: 239,195,904
• Estimated FP32 parameter size: 912.46 MB
• Estimated FP16 parameter size: 456.23 MB
• Serialized checkpoint directory size: 917.11 MB

Browser deployment uses quantized ONNX artifacts rather than the original PyTorch checkpoint.

ONNX export

The model was exported using the image-to-text-with-past task and ONNX opset 18.

The browser export produced:

• encoder_model.onnx
• decoder_model.onnx
• decoder_with_past_model.onnx
• decoder_model_merged.onnx

The files were then processed through dynamic ONNX Runtime quantization.

Browser model files

This repository contains the following quantized ONNX files:

onnx/
├── encoder_model_quantized.onnx
├── decoder_model_quantized.onnx
├── decoder_with_past_model_quantized.onnx
└── decoder_model_merged_quantized.onnx