|
--- |
|
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 |
|
|
|
```python |
|
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) |
|
``` |
|
|
