Skip to content

v1.17.0: Improved ONNX support & many bugfixes
Compare
Choose a tag to compare
@fxmarty fxmarty released this 16 Feb 09:22
· 57 commits to main since this release
v1.17.0
bf7e98e

ONNX export from nn.Module

A function is exposed to programmatically export any nn.Module (e.g. models coming from Transformers, but modified). This is useful in case you need to do some modifications on models loaded from the Hub before exporting. Example:

from transformers import AutoModelForImageClassification
from optimum.exporters.onnx import onnx_export_from_model

model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224")

# Here one could do any modification on the model before the export.
onnx_export_from_model(model, output="vit_onnx")

ONNX export with static shapes

The Optimum ONNX export CLI allows to disable dynamic shape for inputs/outputs:

 optimum-cli export onnx --model timm/ese_vovnet39b.ra_in1k  out_vov --no-dynamic-axes

This is useful if the exported model is to be consumed by a runtime that does not support dynamic shapes. The static shape can be specified e.g. with --batch_size 1 . See all the shape options in optimum-cli export onnx --help.

BF16 ONNX export

The Optimum ONNX export now supports BF16 export on CPU and GPU. Beware though that ONNX Runtime is most often not able to consume the models as some operation are not implemented in this data type, although the exported models comply with ONNX standard. This is useful if you are developing a runtime that consomes BF16 ONNX models.

Example:

optimum-cli export onnx --model bert-base-uncased --dtype bf16 bert_onnx

ONNX export for news models

You can now export to ONNX table-transformer, bart for text-classification.

Sentence Transformers ONNX export

Timm models support with ONNX Runtime

Timm models can now be run through ONNX Runtime with the class ORTModelForImageClassification:

from urllib.request import urlopen

import timm
import torch
from PIL import Image

from optimum.onnxruntime import ORTModelForImageClassification

# Export the model to ONNX under the hood with export=True.
model = ORTModelForImageClassification.from_pretrained("timm/resnext101_64x4d.c1_in1k", export=True)

# Get model specific transforms (normalization, resize).
data_config = timm.data.resolve_data_config(pretrained_cfg=model.config.pretrained_cfg)
transforms = timm.data.create_transform(**data_config, is_training=False)

img = Image.open(
    urlopen("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png")
)
output = model(transforms(img).unsqueeze(0)).logits
top5_probabilities, top5_class_indices = torch.topk(torch.softmax(output, dim=1) * 100, k=5)

Other changes and bugfixes

New Contributors

Full Changelog: v1.16.0...v1.17.0

Contributors

soulteary, dwyatte, and 16 other contributors
Assets 2