OCamlでLightGBMを動かして機械学習してみた

長女と2人ポケパーク、世代を超えて愛されるものを作れるってすごいよね

はじめに

こんにちは、エムスリー株式会社 業務執行役員 VPoE 兼 基盤チーム チームリーダーの河合(@vaaaaanquish)です。

この記事は「基盤開発チーム ブログリレー3日目」の記事です。

先日、基盤チームで「うちもブログリレーやろう！」と盛り上がりまして、「いいね！私も書こうかな！」と気軽に思っていたのですが、何故か1日目と2日目の記事がOCamlの記事になっており「それをやられたらビッグウェーブに乗るしかないが…？」ということで私もOCamlで機械学習をやりました。

つまり本記事は、気合いと根性のやってみた記事になります。

LightGBMとは

LightGBMは、勾配ブースティングというアルゴリズムを実装したフレームワークです。

2016年8月にMicrosoftが公開して以来、汎化性能の高さやライブラリとしての扱いやすさから人気を伸ばし、今尚Kaggleの表形式コンペでは「最も定番のフレームワーク」と言っても差し支えないほど使われています。

github.com

本家の中身はC++で、OpenMPでゴリゴリに並列化しつつメモリ書き込みの競合を避けるためにスレッドローカルなメモリ領域を確保したりと、面白い工夫がいくつも実装されている強いライブラリです。

私はLightGBMのRust版のBindingを作った経験があり、LightGBMにcommitする程度にはC APIを概ね把握しているため、OCamlでもFFIさえ掴めれば出来るはずです*1。

github.com

一方OCamlは初心者。一体、どうなっちゃうのー！？

ということで出来ました

出来ました。

github.com

なんとかなったので、実際の挙動はRepositoryを見てもらうとして、ここからは技術的に面白かった部分の話を書いていこうと思います。

技術的に面白いなと思ったところ

ctypes-foreign

調べるとOCamlでC関数を呼ぶ方法はいくつか見つかりました。

今回は ctypes-foreign を使いました。

github.com

ctypes-foreignは、libffiを利用した動的FFIを扱うライブラリで、OCaml側だけでC関数のバインディングを記述できそうだったので「これなら簡単じゃん…！」という事で進めました。

当然痛い目に遭いました。

具体的にはポインタ周りです。 ctypes-foreignでは、C関数名と型シグネチャを渡す書き方をします。 @-> 演算子で引数の型を連鎖させ、returning で戻り値の型を指定する、以下のようなコードです。

open Ctypes
open Foreign

let lgbm_dataset_create_from_mat =
  foreign "LGBM_DatasetCreateFromMat"
    (ptr void @-> int @-> int32_t @-> int32_t @-> int @->
     string @-> ptr void @-> ptr (ptr void) @-> returning int)

LightGBMは機械学習のライブラリなので行列演算が主たる操作になります。 つまり、ポインタのポインタのポインタを多く扱うライブラリでして、頭の中で「今＠->＠->＠->＠->だから次は＠->＠->で……つまり今俺は＠->＠->＠->を触っている……？」という、AIもやってくれないパズルをすることになりました。

ctypes-foreignは、書き手にCのスタブコードを書かせないという意味で凄いライブラリで、型レベル DSLや動的libffiの実装は他のFFIバインディングでも使えるなというアイデアが詰まっています。一方普遍的な事実として、プログラミングにおいて初見の簡単さは実際の簡単さに比例しないという当たり前の事を思い出す、そんな令和の春になりました。

もうぶっちゃけほとんどこのパズルに時間使ったので、ここからはオマケと言っても過言ではないです。

float32/float64問題

LightGBMのC APIでは、機械学習において、入力となる行列(特徴量)がfloat32/float64、回答を表す行列(ラベル)がfloat32です。

Rustではf64とf32が別の型として存在するので自然に区別できますが、OCamlのfloatは常に64bitのdoubleです。

ここでctypesのCArrayを活用しています。

(* 特徴量: OCaml float → C double（自然な変換） *)
let flat = CArray.make double (nrow * ncol) in
Array.iteri (fun i row ->
  Array.iteri (fun j v -> CArray.set flat (i * ncol + j) v) row
) data;

(* ラベル: OCaml float → C float（暗黙の精度切り捨て） *)
let c_label = CArray.make float label_len in
Array.iteri (fun i v -> CArray.set c_label i v) label;

CArray.make doubleとCArray.make float でCの型を指定して、ctypesに適切なメモリレイアウトのバッファを確保させています。 ctypesのランタイム型記述子で作った事で、64bit→32bitの精度切り捨ても自動で行えて、個人的にニヤニヤしたお気に入りの実装です。

メモリ解放問題

C APIで確保したリソースは明示的に解放が必要になるのは当然のことです。 リソース解放においては、OCamlでは Gc.finalise という機能があります。

実装を始めた当初は決定論的に呼ばれるものかという意識をせず「Rustで言う所のDropトレイトみたいなもんかな」という安易な考えで組み込みました。

let from_mat ~data ~label =
  (* ... C APIでハンドルを取得 ... *)
  let t = { handle } in
  Gc.finalise free t;  (* GCがtを回収する時にfreeが呼ばれる *)
  t

ここでRustのDropにはない制約としてGc.finalise の「ファイナライザ内での例外処理」に当たりました。

A finalisation function may raise an exception; in this case the exception will interrupt whatever the program was doing when the function was called.

上記の公式ドキュメントの記載によると、ファイナライザは例外を投げると、その例外はGCのタイミングに依存した無関係なアロケーションポイントで再raise されるため、GCが発生した時点で実行中だった無関係なコードが中断されるみたいです。

そのため free関数でC APIの戻り値を無視するような実装になっています。

let free t =
  ignore (Lightgbm_raw.lgbm_dataset_free t.handle)

これに気付けずメモリが解放されないままLightGBMの学習が進んでしまい、PCがフリーズするけどフリーズするからデバッグ出来ないというやつを久々に味わう事が出来ました*2。

本気で安全にするなら、明示的なclose / with_datasetを提供し、Gc.finalise はリーク防止の最後の保険にするのが良さそうに思いましたが、体力続かずfreeな気持ちで書きました。

char**のマーシャリング

OCamlでFFIを触るにあたって面白かったのが、C APIの char** の事前確保済みバッファへの書き込みでした。

// C API のシグネチャ
int LGBM_BoosterGetFeatureNames(
  BoosterHandle handle,
  int len,           
  int* num_feature_names,  
  size_t buf_len,  // 各バッファの最大長
  size_t* out_buffer_len,
  char** out_strs    // 事前確保した char* の配列
);

C APIがこんな感じの実装になっている時、OCaml側では、この char**を手動で構築する必要があります。

let feature_name t =
  let nf = num_feature t in
  let buf_len = 256 in
  (* nf個のcharバッファを確保 *)
  let buffers = Array.init nf (fun _ -> CArray.make char (buf_len + 1)) in
  (* char* のポインタ配列を構築 *)
  let ptrs = CArray.make (ptr char) nf in
  Array.iteri (fun i buf -> CArray.set ptrs i (CArray.start buf)) buffers;
  (* ... C API呼び出し ... *)
  (* char* → OCaml string への変換 *)
  List.init actual_num (fun i ->
    let p = CArray.get ptrs i in
    coerce (ptr char) string p)

CArray.make char でバッファを確保し、そのポインタをCArray.make (ptr char) に詰め、C関数に渡して、戻ってきたら coerce (ptr char) string でOCaml文字列に変換しています。

Pythonなら ctypes.c_char_p の配列で、RustならCString::into_raw()やCString::from_raw()で管理するだけですが、OCamlのctypesが低レベル寄りなので、その分何が起きているかが明確に見えて面白かったです。

動かしてみる

おまけ程度ですが、作成したLightGBM OCaml Bindingsを動かす時はこんな感じになります。

(* データセットを用意し学習 *)
let dataset = Lightgbm.Dataset.from_mat ~data:train_features ~label:train_labels in
let params = [
  ("num_iterations", "100");
  ("objective", "binary");
  ("metric", "auc");
] in
let booster = Lightgbm.Booster.train dataset ~params in
(* テスト用データセットに対して推論を実行 *)
let result = Lightgbm.Booster.predict booster ~data:test_features in

実行結果は以下の通り。

[LightGBM] [Info] Number of positive: 3716, number of negative: 3284
[LightGBM] [Info] Total Bins 6132
[LightGBM] [Info] Number of data points in the train set: 7000, number of used features: 28

feature importance:
  Column_25: 211.000000
  Column_27: 179.000000
  Column_24: 178.000000
  Column_22: 173.000000
  Column_0: 164.000000
  ...

result: 381 / 500

自作データセットで正解率 76.2%、randomよりは良いそれなりの数値になったので、一旦動いてはいそうです。

おわりに

OCamlでLightGBMを動かす実装をやってみました。

色々試していたらコードはあまり綺麗に出来ませんでしたが、Claude Codeがメモリの中身までは想像してくれないのを体感出来たので、良い技術者と共に歩む必要がまだまだあるなと確信できました。

今回はテックブログ用なのでctypes-foreignを使いましたが、本格的なライブラリとして配布するなら ctypes-cstubsでCスタブを生成し、ビルド時に型やリンクの問題を検出できる形にする方が堅そうです。

OCamlでC APIを叩く事があれば参考にしてもらえれば嬉しいです。

We are hiring !!

エムスリーでは、ソフトウェアエンジニアを募集しています。 新卒・中途採用、カジュアル面談やインターンも募集しています！

SRE、Platform Engineerを絶賛採用しておりまして、今年はクラウドネイティブ会議に始まり、SRE Nextや関数型まつり、Go Conference等にスポンサー、参加予定です。 会場でのお声がけもお待ちしております。

エンジニア採用ページはこちら

jobs.m3.com

カジュアル面談もお気軽にどうぞ

jobs.m3.com

インターンも常時募集しています

open.talentio.com

*1:私がlightgbm-rsを書いた頃はそんな奇特な人はいなかったんですが今調べると何人か挑戦していて感無量です…良かったねLightGBM

*2:ちゃんと考えると「例外のせい」というよりは「OCamlのGCがC側のメモリ消費量を検知できないため、ファイナライザの実行が遅すぎてメモリリーク」な気がします多分