cloverrose's blog

Python, Machine learning, Emacs, CI/CD, Webアプリなど

Spark版xgboostでRank学習できるようにPR投げた

機械学習で有名なxgboostというライブラリに先日PRがマージされた!嬉しい!

[jvm-packages] call setGroup for ranking task by cloverrose · Pull Request #2066 · dmlc/xgboost · GitHub

PRを送った経緯

  • 会社の業務で検索のランキングモデルを作っていて、xgboostへの移行を検討している。(現在進行形)
  • 会社にはHadoopクラスタがあり、ログはHadoop上に乗っている
  • なのでHadoop上でログの整形→学習が完結するとかなりクール
  • Hadoop上で動くxgboostはYARN版とSpark版がある
    • YARN版はHadoopクラスタにライブラリを追加しないといけないなどの理由でSpark版の方を採用
  • xgboost4j-sparkにはRank学習をサポートしていなかった
  • ただxgboost4j-scalaまでたどればsetGroupメソッドが用意されているので、Sparkからメソッドを呼ぶようにすればよいことがわかった

自分がコードの意図をちゃんと理解していない部分が何個かあったためレビューをたくさんしていただいた。

変更自体は9月くらいに実装して、機械学習のコンペで使っていたのですが、ついにPRを投げてそれがマージされたのは嬉しかった。