【Hummingbird, ONNX Runtime】sklearnで学習した機械学習モデルの推論処理を高速化する

こんにちは。

現役エンジニアの”はやぶさ”@Cpp_Learningです。仕事でもプライベートでも機械学習で色々やってます。

今回は機械学習モデルの推論処理を高速化する方法について勉強したので、備忘録も兼ねて本記事を書きます。

CPUマシンで検証しました（GPUによる高速化はしません）

Contents

1 実践する内容と対象読者
2 sklearnによる機械学習モデルの学習
3 機械学習モデルの推論処理を高速化する２つの方法
4 ONNX Runtimeによる高速化
5 Humingbirdによる高速化
6 まとめ -Hummingbird, ONNX Runtimeによる推論高速化-

実践する内容と対象読者

最初に本記事で実践する内容を紹介します。

本記事のストーリー

sklearnで機械学習モデルを学習（学習済みモデル生成）
学習済みモデルを高速処理に適したフォーマットに変換
変換後のモデルを使って推論処理を高速化

機械学習モデルの高速化については、実に様々なテクニックが存在し、高度な専門スキルを保有していないと実践できないケースもあります。

本記事の内容は sklearnで学習・推論を実践した人を対象としますが、それ以上のスキルは求めません。

sklearnの機械学習モデルを手軽に高速化する方法について紹介しますので、興味のある人は続きをどうぞ。

sklearnによる機械学習モデルの学習

今回は機械学習チュートリアルでお馴染みのirisデータセットを使い、ランダムフォレストによる分類モデルを生成します。

Import

まずはimportから

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier

from sklearn.datasets import load_iris

from sklearn.model_selection import train_test_split

from sklearn.ensemble import RandomForestClassifier

データセット

irisデータセットを取得し、訓練用と評価用に分割します。

# Load dataset
iris = load_iris()
X, y = iris.data, iris.target

# Split dataset (train:test = 7:3)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)

# Load dataset

iris = load_iris()

X, y = iris.data, iris.target

# Split dataset (train:test = 7:3)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)

ランダムフォレストの分類モデルを学習

今回はランダムフォレスト（RandomForestClassifier）を採用して学習します。

# Train a model(RandomForest)
sk_model = RandomForestClassifier()  # default
sk_model.fit(X_train, y_train)

# Train a model(RandomForest)

sk_model = RandomForestClassifier() # default

sk_model.fit(X_train, y_train)

※ハイパーパラメータのチューニングはせず、デフォルト設定を採用

sklearn Ver0.22 以降 n_estimators（木の数）のデフォルト値が10から100に変更したそうです。

推論

学習済みモデルを使い、テストデータのラベルなどを推論します。

%%timeit -r 3
# predict
sk_labels = sk_model.predict(X_test)
sk_prob = sk_model.predict_proba(X_test)

%%timeit -r 3

# predict

sk_labels = sk_model.predict(X_test)

sk_prob = sk_model.predict_proba(X_test)

100 loops, best of 3: 14 ms per loop

推論にかかった時間は約14msでした。この処理時間をどこまで短くできるかが、本記事の本題になります。

Jupyter Notebook（Google Colabも可）のマジックコマンド %%timeit で処理時間を計測しました

推論結果

推論結果は以下のコードで確認できます。

print("score", sk_model.score(X_test, y_test))
print("y_test", y_test)
print("sk_labels", sk_labels)
print("sk_prob", sk_prob)

print("score", sk_model.score(X_test, y_test))

print("y_test", y_test)

print("sk_labels", sk_labels)

print("sk_prob", sk_prob)

出力表示は割愛するので、実際に手を動かして確認してみてください。

（学習フェーズ完）

機械学習モデルの推論処理を高速化する２つの方法

以降からは推論処理の高速化にトライします。今回は主に以下の2つを活用していきます。

ONNX や ONNX Runtime については以下の記事で紹介済みなので割愛します。

ONNX RuntimeとYoloV3でリアルタイム物体検出Microsoft社製OSS”ONNX Runtime”の入門から実践まで学べる記事です。ONNXおよびONNX Runtimeの概要から、YoloV3モデルによる物体検出（ソースコード付）まで説明します。深層学習や画像処理に興味のある人にオススメの内容です。...

Humingbird については、公式が以下のように説明しています。

Hummingbird is a library for compiling trained traditional ML models into tensor computations. Hummingbird allows users to seamlessly leverage neural network frameworks (such as PyTorch) to accelerate traditional ML models.

(中略)

Currently, you can use Hummingbird to convert your trained traditional ML models into PyTorch, TorchScript, ONNX, and TVM). Hummingbird supports a variety of ML models and featurizers. These models include scikit-learn Decision Trees and Random Forest, and also LightGBM and XGBoost Classifiers/Regressors.

引用元：Hummingbird｜GitHub

上記の説明がよく分からない人でも、以降で紹介するソースコードを読みながら実践をすれば理解できると思います。

頭で理解できなくても、手を動かすことで理解できるケースがあります。

ONNX Runtimeによる高速化

最初に ONNX Runtime による高速化を実践します。まずは以下のコマンドで各種のインストールをします。

pip install skl2onnx
pip install onnxruntime

以降からコード書いていきます。

Import

まずはimportから

import numpy
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

import onnxruntime as ort

import numpy

from skl2onnx import convert_sklearn

from skl2onnx.common.data_types import FloatTensorType

import onnxruntime as ort

sklearnのモデルをONNXフォーマットに変換

sklearnのモデルをONNXフォーマットに変換するには、sklearn-onnx を使います。

# Convert into ONNX format
initial_type = [('float_input', FloatTensorType([None, 4]))]
onnx_model = convert_sklearn(sk_model, initial_types=initial_type)

# Save onnx model
with open("rf_iris.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

# Convert into ONNX format

initial_type = [('float_input', FloatTensorType([None, 4]))]

onnx_model = convert_sklearn(sk_model, initial_types=initial_type)

# Save onnx model

with open("rf_iris.onnx", "wb") as f:

f.write(onnx_model.SerializeToString())

わずか数行で変換から”rf_iris.onnx”の保存まで完了です。

ONNX（拡張子が .onnx, .pb, .pbtxt）のモデルは Netron で可視化できます（下図参照）。

ONNX Runtimeによる高速な推論

ONNX Runtime を使うことで、高速な推論を実現します。

最初に先ほど保存した”rf_iris.onnx”をロードして、セッションを作成します。

# Create inference session
session = ort.InferenceSession("./rf_iris.onnx")

1 2	# Create inference session session = ort.InferenceSession("./rf_iris.onnx")

次に各種の名前を取得後、ONNX Runtimeによる推論を実行します。

%%timeit -r 3
# Get name
input_name = session.get_inputs()[0].name
label_name = session.get_outputs()[0].name
probabilities = session.get_outputs()[1].name

# Compute the prediction with ONNX Runtime
onnx_pred = session.run(
                            [label_name, probabilities],  # [output info]
                            {input_name: X_test.astype(numpy.float32)}  # {input info}
                        )

# print("onnx_labels", onnx_pred[0])
# print("onnx_prob", onnx_pred[1])

%%timeit -r 3

# Get name

input_name = session.get_inputs()[0].name

label_name = session.get_outputs()[0].name

probabilities = session.get_outputs()[1].name

# Compute the prediction with ONNX Runtime

onnx_pred = session.run(

[label_name, probabilities], # [output info]

{input_name: X_test.astype(numpy.float32)} # {input info}

)

# print("onnx_labels", onnx_pred[0])

# print("onnx_prob", onnx_pred[1])

10000 loops, best of 3: 131 µs per loop

推論にかかった時間は約131μsでした。桁違いの速さですね。

Humingbirdによる高速化

続いて Hummingbird による高速化を実践します。

Import

まずはimportから


from hummingbird.ml import convert

from hummingbird.ml import convert

sklearnモデルをPyTorchモデルに変換

以下のコードでsklearnモデルをPyTorchモデルに変換します。

# Use Hummingbird to convert the sklearn model to PyTorch
hb_model = convert(sk_model, 'torch')

1 2	# Use Hummingbird to convert the sklearn model to PyTorch hb_model = convert(sk_model, 'torch')

※バックエンドがPyTorchになる

PyTorchモデル（Humingbird製）で推論

以下のようにsklearnと同じ使い心地で簡単に推論できます。

%%timeit -r 3
# Run predictions on CPU
hb_labels = hb_model.predict(X_test)
hb_prob = hb_model.predict_proba(X_test)
# print("hb_labels", hb_labels)
# print("hb_prob ", hb_prob)

%%timeit -r 3

# Run predictions on CPU

hb_labels = hb_model.predict(X_test)

hb_prob = hb_model.predict_proba(X_test)

# print("hb_labels", hb_labels)

# print("hb_prob ", hb_prob)

1000 loops, best of 3: 1.83 ms per loop

推論にかかった時間は約1.83msでした。ONNX Runtimeほどではありませんが、高速化できました。

ONNXモデルをPyTorchモデルに変換

Hummingbird を使えば、ONNXモデルをPyTorchモデルに変換することもできます。

# Use Hummingbird to convert the ONNX-ML model to PyTorch.
onnx_torch = convert(onnx_model, 'torch')

1 2	# Use Hummingbird to convert the ONNX-ML model to PyTorch. onnx_torch = convert(onnx_model, 'torch')

推論は以下の通りです。

%%timeit -r 3
onnx_torch_labels = onnx_torch.predict(X_test.astype(numpy.float32))
onnx_torch_prob = onnx_torch.predict_proba(X_test.astype(numpy.float32))
# print("onnx_torch_labels", onnx_torch_labels)
# print("onnx_torch_prob", onnx_torch_prob)

%%timeit -r 3

onnx_torch_labels = onnx_torch.predict(X_test.astype(numpy.float32))

onnx_torch_prob = onnx_torch.predict_proba(X_test.astype(numpy.float32))

# print("onnx_torch_labels", onnx_torch_labels)

# print("onnx_torch_prob", onnx_torch_prob)

1000 loops, best of 3: 1.7 ms per loop

推論にかかった時間は約1.7msでした。変換元のモデルは違えど、変換後はPyTorchモデルなので、先程とほとんど同じ処理時間ですね。

ONNXモデルをONNXモデル（Hummingbird製）に変換

Hummingbirdで以下のような変換もできます。

# Use Hummingbird to convert the ONNX-ML model to ONNX.
hb_onnx = convert(onnx_model, "onnx")

1 2	# Use Hummingbird to convert the ONNX-ML model to ONNX. hb_onnx = convert(onnx_model, "onnx")

推論は以下の通りです。

%%timeit -r 3
hb_onnx_labels = hb_onnx.predict(X_test.astype(numpy.float32))
hb_onnx_prob = hb_onnx.predict_proba(X_test.astype(numpy.float32))
# print("hb_onnx_labels", hb_onnx_labels)
# print("hb_onnx_prob", hb_onnx_prob)

%%timeit -r 3

hb_onnx_labels = hb_onnx.predict(X_test.astype(numpy.float32))

hb_onnx_prob = hb_onnx.predict_proba(X_test.astype(numpy.float32))

# print("hb_onnx_labels", hb_onnx_labels)

# print("hb_onnx_prob", hb_onnx_prob)

100 loops, best of 3: 5.25 ms per loop

推論にかかった時間は約5.25msでした。ちょっと物足りない気もしますが、ONNX Runtimeは使いたくないけど、ONNXモデルを使いたいケースには良いかと。

くるる

そんなケースあるのー？

うーん。分からん（笑）

はやぶさ

ユースケースに心当たりのある人は教えてください

まとめ -Hummingbird, ONNX Runtimeによる推論高速化-

ONNX Runtime, sklearn-onnx および Hummingbird を活用し、sklearnモデルの変換と推論処理の高速化を実践しました。

結果を以下の表にまとめます。

モデル	推論時間
sklearn（オリジナル）	14ms
sklearn to ONNX（sklearn-onnx製）	131μs（ONNX Runtime使用）
sklearn to PyTorch（Humingbird製）	1.83ms
ONNX to PyTorch（Humingbird製）	1.7ms
ONNX to ONNX（Humingbird製）	5.25ms