Chief Research Officerの西鳥羽 (Jiro Nishitoba (@jnishi) | Twitter) です。今回はHugging Face TransformersのモデルのONNX runtimeで実行を試してみました。それにより特にCUDAでの実行では2.4倍近い高速化が行えました。 Hugging Face TransformersのモデルのONNX形式への変換方法から、ONNX runtimeでの実行も含めて紹介したいと思います。

ONNXとは

ONNX とは、機械学習のモデルを表現するOpenなフォーマットです。機械学習を実現するフレームワークは数多くありますが、ONNX形式での出力をサポートしてるものも多数存在 *1 します。 ONNX形式に変換することによってフレームワークがAPIを用意している言語や関数名にとらわれずに共通化できます。また、ONNX形式の実行時に使われる onnxruntime には様々な最適化があり、推論処理の高速化も実現できます。その汎用性と効率性から、推論用のモデルのデプロイに適しています。

ONNX形式への変換

Hugging Face Transformers*2には学習済みモデルをONNX形式に変換するツールがついています*3。ツールを有効にするには以下のようにオプションをつけてインストールします。

pip install transformers[onnx]

もしくは以下のようにしてソースからインストールします。こちらの場合は特にオプションは必要ありません。

git clone git@github.com:huggingface/transformers.git
cd transformers
pip install -e .

学習済みモデルを用意し、下記のコマンドでONNX形式に変換します。この例では models/bert.model をONNX形式に変換しています。

python -m transformers.onnx --model=models/bert.model --feature sequence-classification onnx/

それぞれの引数は以下のようになっています。

--model transformersの学習済みモデルへのパスを指定します。なお、事前学習済みモデルだけでなく、 BertForSentenceClassification クラスなどダウンストリームタスクに対しても用いることができます。なお、 cl-tohoku/bert-large-japanese のような定義済みのモデルにも適用できます。
--features --model で指定したモデルがどのタスク向けに学習されたものかを指定できます。主なものは以下の通りです*4。文書分類の時は sequence-classification を指定します。
- default
- causal-lm
- masked-lm
- seq2seq-lm
- sequence-classification
- token-classification
- multiple-choice
- question-answering
最後の引数は出力先ディレクトリです

実行に成功すると onnx/model.onnx が出力されます。

ONNX形式に変換したモデルを用いて推論する

onnxruntimeをインストールします。

pip install onnxruntime

モデルの読み込みは以下のようになります。

session = InferenceSession("onnx/model.onnx")

InferenceSession の引数にはONNX形式のモデルへのパスを指定します。

モデルの推論方法は以下のようになります。

 result = session.run(output_names=["logits"], input_feed=dict(eval_dataset[0]))

session.run の output_name には出力に受け取りたい項目を記述します。sequence-classificationの場合は "logits" となります。ほかの場合はONNXに対応しているクラスのconfigを読み込み、その outputs.keys() を表示すると一覧が見られます。

from transformers.models.distilbert import DistilBertConfig, DistilBertOnnxConfig

config = DistilBertConfig()
onnx_config = DistilBertOnnxConfig(config)
print(list(onnx_config.outputs.keys()))
["last_hidden_state"]

推論用データの読み込み、CPUでの推論を通して行うサンプルプログラムが以下になります。

import argparse
import csv

from onnxruntime import InferenceSession
from transformers import AutoTokenizer

def load_csv(filename):
    with open(filename) as f:
        reader = csv.DictReader(f)
        return [row for row in reader]

def main(args):
    # 前処理を行う。ここではrun_glue.py互換の`label`, `sentence1`というヘッダを持つCSVを入力とする
    raw_datasets = load_csv(args.evalfile)

    # Tokenizerの初期化。元のモデルで使用していたTokenizerの設定およびモデルを読み込む
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        use_fast=True,
    )

    # ONNX形式のモデルで推論する際は return_typeを "np" としてnumpy形式のテンソルを返すようにする
    eval_dataset = [tokenizer(raw["sentence1"],
                              padding=True,
                              max_length=128,
                              truncation=True,
                              return_tensors="np") for raw in raw_datasets]
    # ONNX形式のモデルから推論用モデルを作成
    session = InferenceSession(args.onnx_path)

    # ONNX形式のモデルから推論。文書分類なのでoutput_namesには"logit"を指定する。
    result = session.run(output_names=["logits"], input_feed=dict(eval_dataset[0]))

    print(result)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='TensorRT and ONNX inference')
    parser.add_argument("evalfile", help="evaluation CSV file")
    parser.add_argument("model_path", help="Tokenizer model path")
    parser.add_argument("onnx_path", help="ONNX model path")
    args = parser.parse_args()
    main(args)

simple_infer.py として保存し、以下のようにして、第一引数に推論用ファイルのパス、第二引数にTokenizerへのパス、第三引数にONNX形式のモデルファイルへのパスを指定すると実行できます。

python simple_infer.py evaluation.csv cl-tohoku/bert-large-japanese onnx/model.onnx

ONNX形式のモデルからGPUでの推論

ONNX形式のモデルからGPUでの推論を行う場合、まずGPU用のランタイムのインストールを行います。

pip install onnxruntime-gpu

そのあとに InferenceSession のインスタンス作成時に以下のように引数を与えます。

# ONNX形式のモデルから推論用モデルを作成
session = InferenceSession(args.onnx_path, providers=["CUDAExecutionProvider"])

なお、onnxruntime-gpu をインストールした後は providers の指定が必須になります。CPUで実行する際には "CPUExecutionProvider" を渡す必要があります。

実験

文書分類において推論時間の比較を行いました。事前学習済みモデルとして cl-tohoku/bert-large-japanese を利用し、livedoorニュースコーパスの記事を学習データ:評価データを4:1に分割し、カテゴリ推定でファインチューニングを行いました。比較内容としてはHugging Face Transformersのモデルのまま推論した時とONNX形式に変換し、ONNX runtimeを利用して推論した場合の1サンプルあたりの推論時間を比較しています。推論は評価データ全件に対して行い、その平均時間を推論時間としています。バッチサイズは1としています。

f:id:Christopher-727:20220225095510p:plain — transformers(pytorch)のモデルでの推論時間とONNXのruntimeによる推論時間の比較。

CPUでの実行時間は0.1697秒から0.1402秒と約1.2倍の高速化、CUDAでの実行は0.01594秒から0.006671秒と約2.4倍の高速化となりました。

まとめ

Hugging Face Transformersのモデルの推論をONNX runtimeで実行することにより高速化できました。また、そのための手順としてONNX形式の変換およびONNX runtimeでの実行方法のご紹介を行いました。ONNX形式への変換はHugging Face Transformersがツールを提供しているため、容易に行うことができ、割と手軽に試せるようになったと思います。Deep Learningのモデルはどんどん大きくなっていますが、同時に推論の高速化も進んできて応用しやすくなってきていると感じました。

*1:例えば https://onnx.ai/supported-tools.html に挙げられています。

*2:今回は 4.17.0.dev0 で実験しています。

*3:詳細は https://huggingface.co/docs/transformers/serialization を参照してください。

*4:詳細は https://huggingface.co/docs/transformers/serialization#selecting-features-for-different-model-topologies を参照してください。