決定木アンサンブルモデルのデプロイを容易にするライブラリtreeliteの紹介(1)

Pocket

はじめまして。技術開発部の安井です。現在は弊社で運営しているDSPにおける広告効果の分析や機械学習モデルの構築/運用を行なっています。

DSPと機械学習

詳細なDSPに関する説明は省きますが、DSPではSSPから受け取った広告リクエストに対して以下の事項を決定して広告オークションにおける入札応答を行わなければいけません。

  • 入札可能な広告の内、どの広告で入札を行うのか
  • いくらで入札を行うのか

入札する広告、金額を決定にはクリック率、コンバージョン率等の広告リクエストに対する実際のパフォーマンスの値をを活用しています。これらの広告パフォーマンスを入札サーバ上で予測するために弊社では機械学習モデルを活用しています。しかし、DSPにおける機械学習モデルの活用には以下のようなハードルが存在します。

  • 入札サーバを構築しているプログラムと機械学習モデルを学習しているプログラムでは使用しているプログラミング言語が異なる
  • SSPへの入札応答にはシビアな時間制約が存在する

これらを解決するための方法として今回はtreeliteというライブラリをご紹介します。

treelite

treeliteとは近年巷でよく使われているXGBoostやLightGBM等のGBDTの実装ライブラリやscikit-learnに実装されているRandomForestRegressor(Classifier)、GradientBoostingRegressor(Classifier)等のいわゆる決定木をアンサンブルしているモデルのデプロイを容易にするためのライブラリになっています。加えて、モデルが直接Cコードとして吐き出されるため、推論時間の高速化も期待できる実装となっています。開発は以下のリポジトリで行われています。

dmlc/treelite

Referenceはこちら

Treelite: toolbox for decision tree deployment

treeliteを用いたモデルのデプロイのオプションとして次の3つのオプションが示されています。

1.デプロイターゲットのマシンにtreeliteをインストールする

こちらはデプロイ先のマシンにtreeliteをインストールしておくシンプルな方法です。 ゆえにデプロイ先のマシンにpythonがインストールされている必要があります。

2.デプロイターゲットのマシンにランタイムパッケージと一緒にデプロイする

treeliteではランタイムパッケージをアウトプットすることができます。 このランタイムパッケージをデプロイプロセスで同時にターゲットマシンにデプロイすることで、 モデルを使用できるようにする方法です。 こちらもデプロイ先のマシンにpythonがインストールされている必要があります。

3.モデルとなるCコードのみデプロイする

こちらはCコードとしてアウトプットされたモデルのみをターゲットマシンに送る方法です。 Cコードに変換されたモデルは以下のインタフェースから予測値を取得できます。

dataが予測を行うときに必要となる特徴量の集合で、pred_marginはモデルがアウトプットした値をシグモイド関数に通すか通さないかというフラグとなっています(2値分類タスクの場合)。union Entryに関しては以下のように定義されており、特徴量が存在しない場合はmissingに-1が設定されており、特徴量が存在する場合はfvalueに値が格納されている状態を定義します。qvalueに関しては現状使われていないように見受けられます。

golangでのtreeliteモデルの使用方法

上記3種類のデプロイ方法のうち、3番目のCコードのみをデプロイする方法であればpredict関数とのつなぎ込みの部分さえ実装すればpython以外の言語 でもアンサンブルモデルを使用できます。弊社の入札サーバではgolangが採用されており、cgoを使うことでアンサンブルモデルのつなぎ込みを実現することができます。今回はLightGBMのモデルをtreeliteのモデルに変換し、golangから参照する部分までご紹介します。まずは、LightGBMのモデルを読み込みtreeliteのモデルをアウトプットする部分です。LightGBMのモデルはlgb.modelというファイルに保存されていることとします。

5行目のloadでLightGBMのモデルを読み込み6行目で共有ライブラリに変換を行います。その後、7行目で変換した共有ライブラリから必要なヘッダーとCコード、メタ情報を抽出し、zipファイルにまとめます。デフォルトだとzipファイルには以下のファイルがまとめられています。

header.hmain.cが実際のモデル部分でそれらをコンパイルフローがMakefileにまとまっています。recipe.jsonにはモデルコードのメタ情報が格納されています。それではこれらのモデルをgolangから参照できるようにしましょう。コードはこちらとなります。

ディレクトリ構成は以下のようになっていることを想定しており、こちらのコードはmain.goに記述されています。

こちらのコードは16次元のベクトルを受け取り0~1までの確率値を返すモデルを想定しています。予測を行う関数はfunc predictTreelite(fvals []float32) float32ですがまず最初に受け取ったベクトルをバイナリ列の配列に変換しています。これはunion Entryを表現するための型がgolang側に存在しないためこのような形を取っています。今回はベクトルの中にmissingが存在しない想定で実装がされていますが、missingが存在しうる場合は-1のバイナリ列を代わりに入れてあげるようにします。このようにしてバイナリ列の配列に変換されたベクトルをモデルファイルのpredictに渡してあげることでモデルの予測値を獲得します。実際にビルドして実行すると以下のように0~1までの確率値が出力されます。

さいごに

今回はtreeliteでLightGBMで学習されたモデルをtreeliteのモデルに変換して、golangから参照するところまで紹介しました。次回はtreeliteのモデルとLightGBMのモデルとの予測速度比較、golangから参照したtreeliteモデルとgolangで記述されたアンサンブルモデル予測ライブラリleavesとの予測速度比較を行いたいと思います。

treeliteのドキュメント

SQS + Lambdaでs2sの基盤を作った話
アドサーバの実装にGo言語を用いるメリット
決定木アンサンブルモデルのデプロイを容易にするライブラリtreeliteの紹介(2)