Hugging Face OptimumでTransformersをONNXに変換

Hugging Face は、Transformers モデルを Open Neural Network Exchange (ONNX) 形式に変換するための 3 つの異なるパスを提供しており、ユーザーは細かな制御とハイレベルな抽象化の間で選択できます。最も簡略化された方法は Optimum ライブラリを使用することで、変換プロセスを自動化しつつ Hugging Face パイプラインとの互換性を保ちます。

Hugging Face Optimumによるハイレベル変換

Optimum ライブラリは ORTModelForXxx クラスを通じて ONNX 変換を行う、最もユーザーフレンドリーな方法を提供します。from_pretrained() メソッドで from_transformers=True フラグを設定することで、Optimum は自動的に標準的な Transformers モデルをロードし、内部で transformers.onnx パッケージを使用して ONNX に変換します。

Implementation Example:

from optimum.onnxruntime import ORTModelForSequenceClassification

model = ORTModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english", from_transformers=True)

Optimum で変換されたモデルは、予測にすぐに使用できるほか、Hugging Face パイプラインに直接組み込むこともできます。

transformers.onnxによるミッドレベル変換

transformers.onnx モジュールは設定オブジェクトを利用することで変換プロセスを簡素化し、dynamic_axes のような複雑なパラメータをユーザーが手動で定義する必要をなくします。

Implementation Example:

from pathlib import Path
import transformers
from transformers.onnx import FeaturesManager
from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification

model_id = "distilbert-base-uncased-sst-2-english"
feature = "sequence-classification"
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)

model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(model, feature=feature)

onnx_inputs,
    preprocessor=tokenizer,
    model=model,
    config=onnx_config,
    opset=13,
    output=Path("trfs-model.onnx")
)

torch.onnxによるローレベル変換

torch.onnx API は最も細かな制御を提供しますが、input_names、output_names、dynamic_axes など、いくつかのパラメータを手動で指定する必要があります。

Implementation Example:

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_id = "distilbert-base-uncased-finetuned-sst-2-english"
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
dummy_model_input = tokenizer("This is a sample", return_tensors="pt")

torch.onnx.export(
    model, 
    tuple(dummy_model_input.values()),
    f="torch-model.onnx",  
    input_names=['input_ids', 'attention_mask'], 
    output_names=['logits'], 
    dynamic_axes={'input_ids': {0: 'batch_size', 1: 'sequence'}, 
                  'attention_mask': {0: 'batch_size', 1: 'sequence'}, 
                  'logits': {0: 'batch_size', 1: 'sequence'}}, 
    do_constant_folding=True, 
    opset_version=13,
)

Sources