ONNXを使って推論速度を高速にしてみる - BASEプロダクトチームブログ

f:id:HifiroleLorum:20191213172855p:plain

この記事はBASE Advent Calendar 2019の15日目の記事です。 devblog.thebase.in DataStrategyの齋藤(@pigooosuke)が担当します。

ONNXの概要

Open Neural Network Exchange(ONNX)とは、機械学習モデルを表現するフォーマット形式のことです。ONNXを活用すると、PyTorch, Tensorflow, Scikit-learnなどの各種フレームワークで学習したモデルを別のフレームワークで読み込めるようになり、学習済みモデルの管理/運用が楽になります。今回の記事では、よく利用されているLightGBMモデルからONNXへの出力方法の確認と、ONNXの推論を行う実行エンジンであるONNX Runtime上での推論速度の改善がどれほどなのかを検証していきたいと思います。

https://onnx.ai

学習モデルの用意

今回は、KaggleのTitanicデータを使用して、binary classificationの予測モデルを作成します。

Dataset: https://www.kaggle.com/c/titanic

import pandas as pd
from sklearn.model_selection import train_test_split
import lightgbm as lgb

data = pd.read_csv("path/train.csv")
y = data['Survived']
X = data.drop(['Survived', 'PassengerId', 'Name', 'Ticket', 'Cabin'], axis=1)

# カテゴリー変数をbooleanに展開
# 現在、LightGBMのカテゴリー変数を直接ONNXに変換することが出来ないため
category_cols= X.select_dtypes('O').columns.tolist()
X = pd.get_dummies(X, columns=category_cols, drop_first=True, dtype=bool)

X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.1, random_state=2019)
# training
train_data = lgb.Dataset(X_train, label=y_train)
valid_data = lgb.Dataset(X_valid, label=y_valid)
train_params = {
    'task': 'train',
    'boosting_type': 'gbdt',
    'objective': 'binary',
    'num_leaves': 28,
    'learning_rate': 0.01,
    'verbose': 0,
}

gbm = lgb.train(
    train_set=train_data,
    params=train_params,
    num_boost_round=1000,
    valid_sets=[train_data, valid_data],
    early_stopping_rounds=10,
    verbose_eval=10
)

# Training until validation scores don't improve for 10 rounds
# [10] training's binary_logloss: 0.625936 valid_1's binary_logloss: 0.582429
# [20] training's binary_logloss: 0.588612 valid_1's binary_logloss: 0.550251
# ...
# [240]    training's binary_logloss: 0.331639 valid_1's binary_logloss: 0.346977
# [250]    training's binary_logloss: 0.327381 valid_1's binary_logloss: 0.346515
# Early stopping, best iteration is:
# [248]    training's binary_logloss: 0.328348 valid_1's binary_logloss: 0.346271

かなり雑ですが、モデルの用意が出来ました。

# 型の確認
X.info()
# <class 'pandas.core.frame.DataFrame'>
# RangeIndex: 891 entries, 0 to 890
# Data columns (total 8 columns):
# Pclass        891 non-null int64
# Age           714 non-null float64
# SibSp         891 non-null int64
# Parch         891 non-null int64
# Fare          891 non-null float64
# Sex_male      891 non-null bool
# Embarked_Q    891 non-null bool
# Embarked_S    891 non-null bool
# dtypes: bool(3), float64(2), int64(3)
# memory usage: 37.5 KB

# データの確認
X.head()
#  Pclass  Age SibSp   Parch   Fare    Sex_male    Embarked_Q  Embarked_S
# 0    3   22.0    1   0   7.2500  True    False   True
# 1    1   38.0    1   0   71.2833 False   False   False
# 2    3   26.0    0   0   7.9250  False   False   True
# 3    1   35.0    1   0   53.1000 False   False   True
# 4    3   35.0    0   0   8.0500  True    False   True

ONNX変換

ONNXに変換するためには、事前にinputの型を定義する必要があります。用意されている型は以下の通りです。

整数型: Int32TensorType, Int64TensorType
真偽型: BooleanTensorType
浮動小数数型: FloatTensorType, DoubleTensorType
文字列型: StringTensorType
辞書型: DictionaryType
配列型: SequenceType

今回は、全てnumpyのfloat32でinputを受け付けるようにします。
この設定は活用している学習モデルなどによって変わってきます。
例えば、scikit-learnのPipelineを活用して、テキスト入力をtfidfで変換する処理などを含めてONNX化したい場合は、 inputにStringTensorTypeを設定する必要があります。

参考URL: http://onnx.ai/sklearn-onnx/auto_examples/plot_tfidfvectorizer.html#tfidfvectorizer-with-onnx

LightGBMをONNXに変換するためにonnxmltoolsが必要になるので、事前にライブラリをインストールします。 https://github.com/onnx/onnxmltools

import onnxmltools
from onnxmltools.convert.common.data_types import FloatTensorType, BooleanTensorType, Int32TensorType, DoubleTensorType, Int64TensorType

# 入力の型定義
initial_types = [['inputs', FloatTensorType([None, len(X.columns)])]]
# LightGBM to ONNX
onnx_model = onnxmltools.convert_lightgbm(gbm, initial_types=initial_types)
# save
onnxmltools.utils.save_model(onnx_model, "lgb.onnx")
# モデルをvizualize可能
onnxmltools.utils.visualize_model(onnx_model)

inputsは、入力のラベル名です。入力のshapeは[None, 特徴量数]のFloatTensorを指定しています。

参考までに、LightGBMのclassifierのモデルは下図のような構成になっています。(visualize_modelで生成)

入力値を決定木を通じて、予測ラベルと予測確度を出力しています。

推論

ONNX用の実行環境として、Microsoftが出しているonnxruntimeを使います。こちらもインストールします。 https://github.com/microsoft/onnxruntime

import onnxruntime

session = onnxruntime.InferenceSession("lgb.onnx")
# 入力のラベル名の確認
print("input:")
for session_input in session.get_inputs():
    print(session_input.name, session_input.shape)
# 出力のラベル名の確認
print("output:")
for session_output in session.get_outputs():
    print(session_output.name, session_output.shape)

# vizualizeした図と一致
# input:
# inputs [None, 8]
# output:
# label [None]
# probabilities []

# 推論実行
preds = session.run(["probabilities"], {"inputs": X_train.values[0].astype("float32").reshape(1, -1)})
print(preds)
# [[{0: 0.0961046814918518, 1: 0.9038953185081482}]]

# LightGBMの予測
preds = gbm.predict(X_train.values[0].reshape(1, -1))
print(preds)
# array([0.90389532])

第1引数に出力ラベル名（今回はprobabilitiesのみを出力）。
第2引数に入力ラベル名と値をセットして推論を実行します。
予測結果もLightGBMの予測とONNXの予測がちゃんと一致していました。

速度計測

# onnx
%%timeit -r 30
for v in X_train.values:
    pred = session.run(["probabilities"], {"inputs": v.astype("float32").reshape(1, -1)})
# 43.3 ms ± 7.86 ms per loop (mean ± std. dev. of 30 runs, 10 loops each)

# lightgbm
%%timeit -r 30
for v in X_train.values:
    pred = gbm.predict(v.reshape(1, -1))
# 84.4 ms ± 8.96 ms per loop (mean ± std. dev. of 30 runs, 10 loops each)

MacOS 10.14.6 Intel Core i5 3.1 GHz
python=3.7.3
numpy=1.15.2
lightgbm=2.3.1
onnx=1.6.0
onnxconverter-common=1.6.0
onnxmltools=1.6.0
onnxruntime=1.0.0

上記の条件で計測したところ、ONNXモデルはpureなLightGBMに比べて約半分ほどの時間で推論が出来ているのが確認できました。
ONNXは途中で型変換を入れているので厳密に平等な比較とは言えませんが、それでも十分早かったです。

モデルファイルサイズ計測

import pickle
with open("lgb.pkl", "wb") as f:
    pickle.dump(gbm, f, protocol=pickle.HIGHEST_PROTOCOL)

!du -h lgb.pkl
# 740K lgb.pkl
!du -h lgb.onnx
# 500K lgb.onnx

モデルファイルサイズに関しても、pickleでの圧縮に比べ、68%まで軽量化することが出来ました。

今回は、LightGBMでの手順を確認しましたが、https://github.com/onnx では、各種フレームワークの対応が次々に進んでいます。
独自カスタムした計算をしていない限り対応出来ると思うので、学習モデル運用でONNXを検討してみてはいかがでしょうか。

まとめ

今回、LightGBMのモデルからONNX形式でモデル出力をする手順の紹介と、ONNX上での推論速度の検証を行いました。
ONNXを利用することで学習フレームワークに依存せず、高速な推論ができる環境を作ることが出来そうですね。

明日は基盤グループの id:tenkomaさんとOwners Growthの id:MiyaMasa です！お楽しみに！