ONNXランタイムは、モバイルやエッジを含むクロスプラットフォームの推論エンジンです。ONNXモデルに変換することで、モデルを様々な環境で動かすことができます。
huggingface/optimumを利用することで、Transformersのモデルを簡単にONNXに変換できます。
本稿はだいたいHugging Faceの公式ドキュメントに書いてあることです。
環境
- Python 3.10.12
- optimum 1.8.8
今回はGoogle Colabを利用しました。予め!pip install optimum[exporters]でインストールしておきます。
変換
rinna社の提供する日本語GPTモデルであるrinna/japanese-gpt2-xsmallをONNXに変換してみましょう。下記のコマンドで、my_onnx_gpt/というディレクトリに変換されたモデルが出力されます。
!optimum-cli export onnx --model rinna/japanese-gpt2-xsmall my_onnx_gpt/
以上です。お疲れ様でした。
念のため確認してみましょう。
!ls my_onnx_gpt/
出力
config.json decoder_with_past_model.onnx spiece.model
decoder_model_merged.onnx generation_config.json tokenizer_config.json
decoder_model.onnx special_tokens_map.json tokenizer.json
どうやら.onnxに変換されているらしいことが分かります。
実行
optimumを利用してONNXモデルを実行するには、普通のTransformersを利用して実行するときとあまり変わりません。Hugging Faceの公式ドキュメントから、適切なモデルクラスを選択して差し替えてください。今回であればORTModelForCausalLMです。
from transformers import AutoTokenizer, pipeline
from optimum.onnxruntime import ORTModelForCausalLM
MODEL_NAME = "my_onnx_gpt"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = ORTModelForCausalLM.from_pretrained(MODEL_NAME)
onnx_gen = pipeline("text-generation", model=model, tokenizer=tokenizer)
gen = onnx_gen("昔々あるところに")
print(gen)
出力※毎回変化します
[{'generated_text': '昔々あるところに突然現れた - Cの C と言えば、C のキャプチャーを撮影したのが本当らしい。 そのカメラは「 500」という名称で、「'}]