ONNX形式のモデルを使ってみる

この記事は、マネフォアドベントカレンダー2021 11日目の投稿です。

メリークリスマス、世界。
今年もあっという間に終わろうとしていますね。
アドベントカレンダー、やってる?
界隈ではもはや定番となったこの「アドベントカレンダー」ですけど、僕は幼少期に「毎日1個チョコが出てくるイベント」と認識していたので、「チョコを詰める(記事を書く)ために夜中までネタ出し・検証・執筆をする」苦労が隠れていたなんて知りませんでしたね(適当)

ところでこのマネフォアドベントカレンダー2021ですが、昨日のluccafortさんの記事は読まれましたか?
僕は

バカバカしいものであっても継続することが最も重要なメッセージ

という言葉がとても好きです。
なにごとも、途中でやめてしまいたくなることってたくさんあると思うんですが、この記事を読んでいたら「やっぱり続けなければ」と改めて思わされました。

そんなluccafortさんのいい話とはうってかわって本日はゴリゴリのMLトークをさせてください。

BERTはONNXで早くなるのか

どうも、CTO室AI推進部@ken11です。

BERTの話をエンジニアブログでなにか書いたかは覚えてないんですが、BERTいいですよね。
huggingface transformersに東北大がBertJapaneseTokenizerをのせてくれたおかげで、NLPモデルとしては”とっつきやすさ”で群を抜いているのではないでしょうか。

そんなBERTですが、推論にそこそこ時間がかかりがちなことが気になったりします。
まあ推論と一言で言っても、実態としては前処理で時間食ってるパターンなんかもそれなりにあるとは思いますが、この「BERTの推論早くならないかな」の検証として、今日はONNX形式のモデルに変換して速度比較というのをやってみたいと思います。
なんと、自然言語処理の鉄板ライブラリことhuggingface transformersは、PyTorchのモデルを気軽にONNX形式に変換できるものまで用意されているのです。

準備

最初に言っておくと、今回はCPU推論を前提としています。

なにはともかく、 onnxruntime を入れないと始まりません。

$ pip install onnxruntime

また、実験に使う適当なモデルを用意します。
今回は僕が雑につくったBERTのNERのモデル(PyTorch)を使ってみます。

$ ls model/
config.json  pytorch_model.bin  special_tokens_map.json  tokenizer_config.json  training_args.bin  vocab.txt

ONNXに変換する

まず、先ほど用意したモデルをONNXのモデルに変換していきます。
公式ではこの辺で説明されている作業です。

が、ちょっと待ってください。
現在公式で推奨されているのは python -m transformers.onnx のコマンドを使ったものです。
こちらですが、現在BERTのモデルでは default しかサポートされておらず、NERタスクの token-classification は対象外となっています。
なのでこのコマンドで変換しても、そのままNERタスクで使える状態とはなりません。

というわけで今回は別のやり方で変換します。
今回はこちらconvert_graph_to_onnx を使っていきます。

from pathlib import Path
from transformers.convert_graph_to_onnx import convert, optimize, quantize

# 変換元のモデルがあるディレクトリを指定します
model = 'model'

output_path = "onnx/converted_model/model.onnx"

# 実はこちらはこちらでpipelineとしてtoken-classificationがサポート外なんですが、変換できちゃうのでそのままいきます
# https://github.com/huggingface/transformers/blob/master/src/transformers/convert_graph_to_onnx.py#L32-L42
convert(pipeline_name="token-classification", framework="pt", model=model, tokenizer=model, output=Path(output_path), opset=12)

# これはオマケで、さらにoptimize/quantizeすると速くなるので試します
optimized_output = optimize(Path(output_path))
quantize(optimized_output)

実行すると以下のモデルができあがります。

$ ls onnx/converted_model/
model-optimized-quantized.onnx  model-optimized.onnx  model.onnx

変換自体はこれだけです。
こんな変換まで用意されてるhuggingface transformersは素晴らしい…

コード中にも書きましたが、一応 token-classification はこの方法のサポート外らしいのでご注意ください。

比較してみる

さて、モデルの準備も出来たので比較してみましょう。

推論について

まずはONNX化したモデルの推論ですが、以下のようにして行います。

from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions


options = SessionOptions()
options.intra_op_num_threads = 1
options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL

ort_session = ort.InferenceSession("onnx/converted_model/model.onnx", options, providers=["CPUExecutionProvider"])
ort_session.disable_fallback()
pred = ort_session.run(None, dict(inputs))

ちなみに、ぱっと見では特に推論結果は変わってないようです。
とはいえ今回は精度については比較していないのであしからず。

比較した結果

速度を見たいだけなので、モデルを作成したときに使用した学習データ(約5000の文章)をひたすら推論させて比べてみたいと思います。

通常のPyTorchでの推論と、上述のONNXでの推論、それぞれでどれくらい差があるのか見てみましょう。
なお、速度計測は普通に time でやります。(よくあるやつ)
計測した時間をarrayに格納しておいたので、ざっとどんな感じか見てみましょう。

countはトータルの推論回数、5343回実行しています。
あとは推論時間の集計ですね。
当然ながら入力した文章の長さにもよると思いますが、平均して0.04秒程度ONNXのほうが速いことがわかりました。
たかが0.04秒、されど0.04秒
0.15秒が0.11秒ということは1.3倍程度速くなっています。
今回の処理ではもともとの推論時間が短いのでインパクトが小さく見えるかもしれませんが、1.3倍速くなると思うと結構なインパクトではないでしょうか。

量子化したものも比較してみた結果

ONNX変換する際に、おまけでquantizedモデルまでつくっていました。( model-optimized-quantized.onnx )
ついでにこの量子化モデルについても比較してみましょう。

先ほどよりさらに0.01秒速くなりましたね。
貢献度としては ONNX化 > 量子化 というのはわりとセオリー通りではないでしょうか。
トータルで平均0.05秒速くなったということは1.5倍ですかね。かなり恩恵が大きいと思います。

ただ、量子化したモデルはさすがに推論結果にも差があると思うので、実際に投入するのであれば精度もきちんと見る必要があります。(今回は割愛)

まとめ

huggingface transformersを使ってPyTorchなBERTモデルをONNX形式に変換して速度比較を行ってみました。
機械学習モデルをサービスで活用するには当然速度の問題にも直面すると思います。
huggingface transformersであればこのように簡単にONNXに変換して利用することができるので、選択肢としても候補に入れやすいのではないかなと感じました。

AI推進部ではこのようにMLモデルの実戦投入にあたって試行錯誤することが好きな方を募集しています。
ご応募お待ちしてます。


【サイトのご案内】
マネーフォワード採用サイト
Wantedly
京都開発拠点

【プロダクトのご紹介】
お金の見える化サービス 『マネーフォワード ME』 iPhone,iPad Android

ビジネス向けバックオフィス向け業務効率化ソリューション 『マネーフォワード クラウド』

おつり貯金アプリ 『しらたま』

お金の悩みを無料で相談 『マネーフォワード お金の相談』

だれでも貯まって増える お金の体質改善サービス 『マネーフォワード おかねせんせい』

金融商品の比較・申し込みサイト 『Money Forward Mall』

くらしの経済メディア 『MONEY PLUS』

Pocket