README.md

metadata

license: apache-2.0

This model is a fork of (facebook/levit-256)[https://huggingface.co/facebook/levit-256], where:

• nn.BatchNorm2d and nn.Conv2d are fused
• nn.BatchNorm1d and nn.Linear are fused

and the optimized model is converted to the onnx format.

How to use

from optimum.onnxruntime.modeling_ort import ORTModelForImageClassification
from transformers import AutoModelForImageClassification

ort_model = ORTModelForImageClassification.from_pretrained("fxmarty/levit-256-onnx")

inp = {"pixel_values": torch.rand(4, 3, 224, 224)}

res = model(**inp)
res_ort = ort_model(**inp)

assert torch.allclose(res.logits, res_ort.logits, atol=1e-4)