MosaicML Composer

Composerは、ニューラルネットワークをより良く、より速く、より安価にトレーニングするためのライブラリです。ニューラルネットワークのトレーニングを加速し、一般化を向上させる最先端の多くの手法を提供し、さまざまな強化機能を簡単に組み合わせるためのオプションのTrainer APIが含まれています。

W&Bは、ML実験のログを記録するための軽量なラッパーを提供します。しかし、自分で両方を組み合わせる必要はありません。Weights & Biasesは、WandBLoggerを介してComposerライブラリに直接組み込まれています。

コード2行でW&Bにログを記録する

from composer import Trainer
from composer.loggers import WandBLogger

wandb_logger = WandBLogger(init_params=init_params)
trainer = Trainer(..., logger=wandb_logger)

ComposerのWandBLoggerを使用する

Composerライブラリには、Trainerと一緒に使用して、指標をWeights and Biasesにログを記録するWandBLoggerクラスがあります。ロガーをインスタンス化してTrainerに渡すだけです。

wandb_logger = WandBLogger()
trainer = Trainer(logger=wandb_logger)

Logger引数

以下は、WandbLoggerでよく使用されるパラメータのいくつかです。全てのリストと説明については、Composerのドキュメントを参照してください。

パラメータ	説明
init_params	wandb.initに渡すパラメータ、例えばwandbのproject、entity、nameやconfigなどです。こちらでwandb.initが受け入れる完全なリストを参照してください。
log_artifacts	wandbにチェックポイントをログするかどうか
log_artifacts_every_n_batches	アーティファクトをアップロードする間隔です。log_artifacts=Trueの場合のみ適用されます。
rank_zero_only	ランクゼロプロセスでのみログするかどうかです。wandbにartifactsをログする場合、すべてのランクでログすることを強くお勧めします。ランク≥1のアーティファクトは保存されず、重要な情報が破棄される可能性があります。例えば、DeepSpeed ZeROを使用する場合、すべてのランクのアーティファクトがなければ、チェックポイントから復元することは不可能です（デフォルト：False）

典型的な使用方法は以下の通りです。

init_params = {"project":"composer", 
               "name":"imagenette_benchmark",
               "config":{"arch":"Resnet50",
                         "use_mixed_precision":True
                         }
               }

wandb_logger = WandBLogger(log_artifacts=True, init_params=init_params)

予測サンプルのログ

Composerのコールバックシステムを使用して、WandBLoggerを介してWeights & Biasesにログするタイミングを制御できます。この例では、検証画像と予測のサンプルをログします。

import wandb
from composer import Callback, State, Logger

class LogPredictions(Callback):

    def __init__(self, num_samples=100, seed=1234):

        super().__init__()

        self.num_samples = num_samples

        self.data = []

        

    def eval_batch_end(self, state: State, logger: Logger):

        """バッチごとに予測を計算し、self.dataに格納する"""

        

        if state.timer.epoch == state.max_duration: #最後のvalエポックに

            if len(self.data) < self.num_samples:

                n = self.num_samples

                x, y = state.batch_pair

                outputs = state.outputs.argmax(-1)

                data = [[wandb.Image(x_i), y_i, y_pred] for x_i, y_i, y_pred in list(zip(x[:n], y[:n], outputs[:n]))]

                self.data += data

            

    def eval_end(self, state: State, logger: Logger):

        "wandb.Tableを作成し、ログに記録する"

        columns = ['image', 'ground truth', 'prediction']

        table = wandb.Table(columns=columns, data=self.data[:self.num_samples])

        wandb.log({'sample_table':table}, step=int(state.timer.batch))         

...



trainer = Trainer(

    ...

    loggers=[WandBLogger()],

    callbacks=[LogPredictions()]

)