LightGBMの進捗をプログレスバーで表示する

pip install tqdm

# str | Noneって書きたいところだけど3.7系でもいけるようにOptionalを使う
from typing import Optional
from collections import OrderedDict
from lightgbm.callback import CallbackEnv
from tqdm.auto import tqdm

class LgbmProgressBarCallback:
    description: Optional[str]
    pbar: tqdm

    def __init__(self, description: Optional[str] = None):
        self.description = description
        self.pbar = tqdm()

    def __call__(self, env: CallbackEnv):

        # 初回だけProgressBarを初期化する
        is_first_iteration: bool = env.iteration == env.begin_iteration

        if is_first_iteration:
            total: int = env.end_iteration - env.begin_iteration
            self.pbar.reset(total=total)
            self.pbar.set_description(self.description, refresh=False)

        # valid_setsの評価結果を更新
        if len(env.evaluation_result_list) > 0:
            # OrderedDictにしないと表示順がバラバラになって若干見にくい
            postfix = OrderedDict(
                [
                    (f"{entry[0]}:{entry[1]}", str(entry[2]))
                    for entry in env.evaluation_result_list
                ]
            )
            self.pbar.set_postfix(ordered_dict=postfix, refresh=False)

        # 進捗を1進める
        self.pbar.update(1)
        self.pbar.refresh()

booster = lgbm.train(
    {**param, "verbosity": -1},
    dataset_train,
    valid_sets=dataset_valid,
    num_boost_round=10,
    callbacks=[
      LgbmProgressBarCallback(description="Model A"),
    ],
)

# evaluation_result_listの例
[
  ("valid_0", "binary_logloss", 0.034567, False),
  ("valid_0", "auc", 0.77, True),
  ("valid_1", "binary_logloss", 0.023456,  False),
  ("valid_1", "auc", 0.83, True),
]