KOMIYA Atsushi

KOMIYA Atsushi

Software engineer in Tokyo.

Apache Spark 上で XGBoost の予測モデルを手軽に扱いたい！

TL;DR: Pure Java 実装な XGBoost 互換の予測専用モジュール xgboost-predictor を基に、Apache Spark 上で お手軽 に XGBoost の予測モデルをロードしたり予測を実現するモジュール xgboost-predictor-spark を作りましたよ、というお話です。

(xgboost-predictor のバージョン 0.2.0 リリースノートを兼ねています)

背景

DMLC が提供する勾配ブースティングツリーの実装 XGBoost では、JVM 環境向けに XGBoost4J なるパッケージが公式提供されています。この XGBoost4J には、Java / Scala 向けのインタフェースだけではなく、 Apache Spark / MLlib の Spark ML API にだいたい準拠したモジュール XGBoost4J-Spark も含まれています。

この XGBoost4J-Spark を使うことで、Spark アプリケーション上で DataFrame-friendly な API を通じて XGBoost の予測モデルを構築したり、またその予測モデルを利用した具体的な予測処理ができるようになります。 ただし、XGBoost4J および XGBoost4J-Spark は Maven Central などで公開されているわけではない 1 ため、実際に Spark アプリケーション上で XGBoost4J-Spark を使おうとすると、

  1. XGBoost 自体をビルドする (gcc などが必要)
  2. XGBoost4J & XGBoost4J-Spark をビルドする

といった作業を別途実施する必要が生じます。なお XGBoost4J はネイティブライブラリを内包しているので、「開発環境でビルドした XGBoost4J / XGBoost4J-Spark を含む Spark アプリケーションの fat jar が、プロダクション環境では動かない〜 😫」みたいなことも状況次第では生じ得ます。そのため、XGBoost4J / XGBoost4J-Spark は現時点ではまだ 手軽に使える 状況ではない、と言えるでしょう。

その一方で、元々オンライン & リアルタイム予測処理のために開発していた xgboost-predictor に対して「(Spark で使うために) Scala の API は提供しないのか？ 2」とか「XGBoost4J-Spark が生成する予測モデルファイルが xgboost-predictor で読み込めないぞ 3」といった 当初の目的に対して斜め上の 要望が挙げられるようになったので、これ以上よくわからん issue が追加される前にいっそのこと Spark ML の API に準拠した DataFrame-friendly なモジュールを作ってしまおう、と思い立って xgboost-predictor-spark を作った… というわけです。

xgboost-predictor-spark が提供する機能

xgboost-predictor-spark は、xgboost-predictor を Scala & Spark ML の API でラップしたモジュールに過ぎないので、基本的には xgboost-predictor に準じた機能を提供しています。より具体的には、以下の機能を提供しています。

予測モデルのロード

libxgboost が生成するモデルフォーマットに加えて、XGBoost4J-Spark が生成するモデルフォーマット (libxgboost のモデルフォーマットに Spark 固有のヘッダが追加されたもの) の両方に対応しています。

予測処理

二値分類、多クラス分類、線形回帰それぞれの予測処理をサポートしています。いずれのタスクにおいても、leaf prediction (各ツリーにおいてたどり着いた葉ノードのインデックスを、予測ラベルの代わりに出力する機能) を利用できます。

また分類タスクにおいては、各クラスへの分類確率や margin の値の出力にも対応しています。

API 構成

二値分類、多クラス分類、線形回帰のそれぞれについて、以下のようにモデルを読み込むためのメソッドおよび予測モデルを表すクラスを提供しています。

  • 二値分類
    • モデルの読み込み: XGBoostBinaryClassification.load() メソッド
    • モデルクラス: XGBoostBinaryClassificationModel クラス
      • ProbabilisticClassificationModel クラスを継承しています
  • 多クラス分類
    • モデルの読み込み: XGBoostMultiClassification.load() メソッド
    • モデルクラス: XGBoostMultiClassificationModel クラス
      • ProbabilisticClassificationModel クラスを継承しています
  • 回帰
    • モデルの読み込み: XGBoostRegression.load() メソッド
    • モデルクラス: XGBoostRegressionModel クラス
      • RegressionModel クラスを継承しています

それぞれのモデルクラスで共通的に提供しているパラメータは以下のとおりです。

  • predictLeaves
    • 具体的なラベルの予測や回帰ではなく、leaf prediction をする (true) か否 (false) かを指定します
    • 分類モデルの場合においてこのパラメータに true が指定されている場合は、 rawPredictionCol や probabilityCol で指定されたカラムへの中間値の出力はされません
    • Leaf prediction の予測結果は、predictionCol に DenseVector の型で出力されます
  • missingValue
    • 欠損値を表す値 (0 や Double.NaN など) を指定します
    • 特徴ベクトルが Spark ML の DenseVector で表現されている場合に、このパラメータで設定された値を利用します

利用方法

Apache Spark および Scala のバージョン

xgboost-predictor-spark は Apache Spark 2.0 以降のバージョンでの利用を前提としています。また、Scala のバージョンは 2.10 系、2.11 系の両方をサポートしています。

依存モジュールの定義

xgboost-predictor-spark はまだ JCenter での公開をしていないため、以下のように resolver の追加が必要になります。

resolvers += Resolver.bintrayRepo("komiya-atsushi", "maven")

libraryDependencies ++= Seq(
  "org.apache.spark" %% "spark-mllib"             % "2.0.2" % "provided",
  "biz.k11i"         %% "xgboost-predictor-spark" % "0.2.0"
)

サンプルコード

xgboost-predictor-spark の README に記載しているコードや、xgboost-predictor-examples のコード をご参照ください。

まとめと今後の予定

xgboost-predictor を単に Spark ML の API でラップしただけではありますが、少なくとも需要はあるっぽいので xgboost-predictor-spark のモジュールを作ってみました。

なお、予測モデルの構築 (学習) 機能は一切存在しないので、Spark 上でアドホックに XGBoost の予測モデルを構築して予測性能を評価する… みたいな使い方はできません。予測モデルは Spark の外で構築しておいて、Spark 上ではバッチ処理的に予測する、といった使い方に限定されるかと思われます。

今後は xgboost-predictor 本体のアップデートに合わせて、こちらも更新していくつもりでいます (裏を返せば、積極的な機能追加は要望がない限りは実施しない、ということです 😝)。

  1. https://github.com/dmlc/xgboost/issues/1807 において、Maven Central にて公開する準備がされているようですが、なかなか進捗が思わしくないようです。 ↩

  2. https://github.com/komiya-atsushi/xgboost-predictor-java/issues/9 ↩

  3. https://github.com/komiya-atsushi/xgboost-predictor-java/issues/15 ↩

You may also enjoy