README.md

---
language:
  - en
library_name: transformers
pipeline_tag: image-to-text
base_model: nlpconnect/vit-gpt2-image-captioning
tags:
  - vision-encoder-decoder
  - image-captioning
  - vision-language
  - onnx
  - transformers.js
  - onnxruntime
  - pytorch
  - vit
  - gpt2
---

# 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:

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