pranjal-pravesh/vit-base-patch16-224-onnx

ViT-Base/16-224 — ONNX

ONNX export of google/vit-base-patch16-224 (Vision Transformer, base, patch size 16, 224×224 input). Single self-contained .onnx file, no external weights.

• Task: image classification (1000 ImageNet-1k classes)
• Input: pixel_values, shape (batch, 3, 224, 224), float32, normalized with mean/std (0.5, 0.5, 0.5)
• Output: logits, shape (batch, 1000), float32
• Quantization: none — weights and activations are full precision (float32). No quantization or pruning has been applied.

Inference

Only onnxruntime, numpy, Pillow, and requests are required — no transformers / PyTorch needed.

import io, json, requests, numpy as np, onnxruntime as ort
from PIL import Image

sess = ort.InferenceSession("vit-base-patch16-224.onnx", providers=["CPUExecutionProvider"])

# 1) Load image
url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/hub/parrots.png"
image = Image.open(io.BytesIO(requests.get(url, timeout=30).content)).convert("RGB")

# 2) Preprocess: resize to 224x224, scale to [0,1], normalize with mean=std=0.5, CHW, add batch
image = image.resize((224, 224), Image.BILINEAR)
arr = np.asarray(image, dtype=np.float32) / 255.0
arr = (arr - 0.5) / 0.5
pixel_values = arr.transpose(2, 0, 1)[None]   # shape (1, 3, 224, 224)

# 3) Run
logits = sess.run(None, {"pixel_values": pixel_values})[0]   # (1, 1000)
pred = int(logits.argmax(-1)[0])

# 4) (Optional) map index -> ImageNet-1k class name
labels = requests.get(
    "https://huggingface.co/datasets/huggingface/label-files/raw/main/imagenet-1k-id2label.json",
    timeout=30,
).json()
print(pred, "->", labels[str(pred)])

Output classes are the standard 1000-way ImageNet-1k label set in the same order as the upstream model.

Conversion details

• Source: google/vit-base-patch16-224 (PyTorch safetensors)
• Exporter: torch.onnx.export (PyTorch 2.11, dynamo path)
• Opset: 17
• Dynamic axes: batch dimension on both pixel_values and logits
• Wrapping: the ViTForImageClassification module is wrapped to return only logits (the HF ImageClassifierOutput is dropped) so the ONNX graph has a single named output
• External data: disabled — weights are embedded in the .onnx file (≈330 MB)

Parity vs. PyTorch

Verified on the HF parrots image (FP32 CPU, ONNX Runtime 1.25):

Tensor	max abs err	max rel err	cosine sim
logits	1.34e-05	8.00e-04	1.0000000000
softmax probs	1.19e-07	1.52e-05	1.0000000000

Top-1 and top-5 predictions match the PyTorch model.

License & attribution

Apache-2.0, inherited from the upstream model. See the original model card for training data, procedure, evaluation, and citations.