この記事は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 です!お楽しみに!