xgboost4j より数千倍速く predict できる Pure Java な XGBoost 互換の予測器を作ってみた

TL;DR

XGBoost で構築した予測モデルを Java から利用したい、それも特徴ベクトルが一つ一つ、任意のタイミングで与えられるような オンライン環境下 で リアルタイムな予測 を実現するために利用したい、という目的を叶えるためのモジュールを作りました。

• Github: komiya-atsushi/xgboost-predictor-java
• Bintray: xgboost-predictor

(XGBoost の凄さとか XGBoost そのものの使い方とか GBDT/GBRT の解説は本エントリにはありませんので、そのような情報を求めている方は他のブログエントリを読まれることをおすすめします。)

xgboost4j という選択肢

Java から XGBoost を利用しようとすると、XGBoostをJavaのwrapperを使用して実行する - TASK NOTES にあるように、DMLC が提供している xgboost4j を利用する手段がすでに存在しています。ただ、この xgboost4j をオンライン予測に適用する場合、下記に挙げるようないくつかの懸念があります。

• XGBoost の Java wrapper でしかないので、オンライン予測の目的で利用をするにはインタフェースがちょっと使いづらい
  • 大量の特徴ベクトルを入力して一括予測するようなバッチ処理に適したインタフェースになっている
• LIBSVM フォーマットじゃないデータを入力するのに手間がかかる
  • 特に疎な特徴ベクトルを DMatrix で表現するのが面倒
• JNI 由来のオーバーヘッドが気になる
  • 予測処理などが C++ で書かれているので高速処理が期待できる一方で、特徴ベクトルを一つ一つ与えて予測させる場合、ネイティブコードの呼び出しにかかるオーバーヘッドが全体のパフォーマンスに大きな影響を与えそう
• どこの Maven repository にもアップロードされていないので、自前で mvn install する必要がある
  • OS X で開発をしている場合、ネイティブライブラリをビルドするのも一苦労 (参考)

Pure Java での予測を実現する

そういうわけで、

• 速度性能がオンラインでの利用に耐えうる水準で
• そこそこ使いやすいインタフェースで
• ネイティブコードを必要としない (= Pure Java な)

XGBoost 互換な予測器を作って jCenter で公開しました。

• Github: komiya-atsushi/xgboost-predictor-java
• Bintray: xgboost-predictor

使い方

README.md にもサンプルコード込みで使い方を書いていますので、合わせてご参照ください。

なお、タイトルでも「予測器」と明言しているとおり、学習 (training) 機能については割りきって一切実装をしていません (モデルの構築をオンラインですることは流石にないと思うので)。 そのため学習データを用意し、XGBoost を CLI で直接実行、もしくは Python や R のラッパーを経由するなどして別途モデルの構築を事前に済ませておく必要があります。

学習済みのモデルが用意できたら、そのモデルを new Predictor("/path/to/model-file") としてロードします。 予測をするには Predictor#predict() のメソッドを呼び出します。

予測の際の入力となる特徴ベクトルは、FVec インタフェースのオブジェクトとして表現する必要があります。 基本はお手持ちのデータの形式に合わせて FVec インタフェースを実装したクラスを用意していただくことになりますが、配列や Map でデータが表現されている場合には、以下のユーティリティメソッドを利用することができます。

• double の配列で表現された密な特徴ベクトルで → FVec.Transformer#fromArray() で FVec オブジェクトに変換できます
• 特徴量のインデックスと値の対を表現した Map オブジェクト → FVec.Transformer#fromMap() で FVec オブジェクトに変換できます

また GBDT/GBRT のモデルを特徴ベクトルの変換に利用することを目的に、(分類・回帰の結果を出力するのではなく) GBDT/GBRT の各ツリーにおいて辿り着いたリーフのノード番号を出力する Predictor#predictLeaf() メソッドも用意されています。

ベンチマーク

元々このモジュールを作った動機の一つとして、オンライン利用に耐えうる速度性能で予測をしたい、という目的があったので xgboost4j と合わせてベンチマークをとってみました。 ベンチマーク計測の詳細は こちら のプログラムにあるとおりです。

機能	xgboost-predictor	xgboost4j	性能比
モデルの読み込み	49017.60 ops/s	39669.36 ops/s	1.24
単一の特徴ベクトルでの予測	6016955.46 ops/s	1018.01 ops/s	5910.48
複数の特徴ベクトルの予測	44985.71 ops/s	5.04 ops/s	8931.47
リーフノードの出力	11115853.34 ops/s	1076.54 ops/s	10325.53

結果は上記のとおり、

• 予測処理は 1ms 以下で処理できている
• xgboost4j と比較して約 6,000 倍近い速度性能がでている (要は xgboost4j よりも十分速い)

となります。Pure Java にしただけで xgboost4j との性能差がこんなに出るものなのか… とちょっと不思議ではありますが、今回のベンチマークの計測に利用したテストデータは人工的に生成されたものなので、その影響があるのかもしれません。 そのため、実世界のデータを与えた場合にはまた違った結果になる可能性がありますのでご注意ください。

制限

現時点の xgboost-predictor は、モデルとしては “gbtree” のみをサポートしています。また目的関数は “binary” および “multi” のみのサポートとなります。 (“gblinear” や他の目的関数のサポートについては、必要そうであれば対応する予定です。)