Weights & Biases で K-fold 交差検証のロギング

素直に書くと、こうだけど…

for fold, (train_indices, validation_indices) in enumerate(kf.split(dataset)):
    
    for epoch in range(1, num_epochs + 1):

        # …学習…
    
        wandb.log({
            f'loss/train/fold-{fold}': train_loss,
            f'acc/train/fold-{fold}': train_acc,
            f'loss/validation/fold-{fold}': validation_loss,
            f'acc/validation/fold-{fold}': validation_acc
        })

グラフの横軸（“Step”）は wandb.log が呼び出されるごとにインクリメントされるので、時間軸方向に各 fold の学習が並んでしまう。これだと、特定の fold だけ変な学習をしてるとか、そういう事態に気付きにくくて、不便。エポックを横軸にしたい。

どうするか？まずは wandb.log でエポックもログを取るようにする。

for fold, (train_indices, validation_indices) in enumerate(kf.split(dataset)):
    
    for epoch in range(1, num_epochs + 1):

        # …学習…
    
        wandb.log({
            f'loss/train/fold-{fold}': train_loss,
            f'acc/train/fold-{fold}': train_acc,
            f'loss/validation/fold-{fold}': validation_loss,
            f'acc/validation/fold-{fold}': validation_acc,
            'epoch': epoch
        })

そして、（Web のインターフェース上で）グラフの横軸を Step から epoch に変える。

時間軸ではなくエポック軸でプロットできて便利。また、wandb.log には step という引数があるので、これにエポックの番号を渡せばいいのでは？と思ったが…

wandb.log({
    f'loss/train/fold-{fold}': train_loss,
    f'acc/train/fold-{fold}': train_acc,
    f'loss/validation/fold-{fold}': validation_loss,
    f'acc/validation/fold-{fold}': validation_acc
}, step=epoch)

どうやら既存の Step に値を追記することはできないらしいので、やはり自分でエポック軸を作ってやるのがベストプラクティスっぽい。