TensorBoardでXGBoostの学習結果を可視化する

2022.05.16

データアナリティクス事業本部の鈴木です。

TensorBoardはTensorFlowが提供する可視化ツールです。TensorFlow以外のライブラリからでも、ライブラリ自体の機能や、連携用に開発されたライブラリを使うことで、別の機械学習ライブラリで作成したモデルでも情報を可視化することが可能です。

今回はXGBoostの学習状況をTensorBoardで可視化する方法がないか探していたところ、以下のIssueでtensorboardXを使った方法が紹介されていたので試してみました。

準備

前提

以下の条件で検証を行いました。

  • コンテナ:jupyter/datascience-notebook
  • xgboost: 1.6.1
  • tensorflow: 2.8.0
  • tensorboardX: 2.5

やってみる

Jupyter Notebookでコードを実行して検証してみました。コード例は以下のレポジトリに格納しています。

docker-compose upでコンテナを起動し、Jupyter Notebookにアクセスすることで、コードを実行することが可能です。以降はポイントと思った箇所を抜粋して記載します。

Callbackの作成

今回使ってみた、最新のXGBoost1.6はPythonパッケージの場合、新しいcallbackインターフェースを使う必要がありました。古い書き方をしていると、Old style callback was removed in version 1.6.というようなエラーが出てしまいました。

XGBoostのDemo for using and defining callback functionsと、tensorboardXのadd_scalarのガイドを参考にCallback用のクラスを作成しました。

今回は以下のように、after_iterationメソッドを実装しました。イテレーションごとに訓練およびテストの最新の損失関数の値を記録します。

class TensorBoardCallback(xgb.callback.TrainingCallback):
    '''
    Callback for add loss to summary.
    '''
    def __init__(self):
        self.writer = SummaryWriter()

    def _get_key(self, data, metric):
        return f'{data}-{metric}'

    def after_iteration(self, model, epoch, evals_log):
        '''Add scalar loss data to summary.'''
        self.writer.add_scalar("train_loss", evals_log["train"]["rmse"][-1], epoch)
        self.writer.add_scalar("test_loss", evals_log["test"]["rmse"][-1], epoch)

TensorBoardで結果を確認する

Jupyter Notebook上から以下のコマンドを実行することで、TensorBoardを起動しました。

%tensorboard --logdir ./runs --load_fast=false --bind_all --port 6006

今回はdatascience-notebookイメージを使ってコンテナで環境を作っていますが、ポイントとして、あらかじめ6006番ポートが使えるようにしています。

docker-compose.yml

version: "3"
services:
    datascience-notebook:
        image: jupyter/datascience-notebook:latest
        volumes: 
            - ./:/home/jovyan/work
        ports:
            - 8888:8888
            - 6006:6006
        container_name: tensorboard-sample-container
        environment:
            - JUPYTER_ENABLE_LAB=yes

何回か学習させてみると、TensorBoardからは以下のように可視化できることができました。

TensorBoardからの可視化

最後に

tensorboardXとXGboostの新しいcallbackの仕組みを使って、学習時の損失関数をTensorBoardで確認してみました。新しいcallbackのインターフェースの利用例としても、参考になれば幸いです。

参考