【機械学習】PyTorch Lightningのログを実験管理ツールwandbで記録する

PyTorch Lightningでwandbを使って実験管理をしたい!しかしドキュメント見てもよくわからなかった…という方に向けて、こちらに残しておこうと思います。

今回も、前に書いたときと同様に公式ドキュメントにあったCoolSystemの実装コードをもとに改修していきます。

※PyTorch Lightningのgithubの内容が最近変わったようで、この実装コードが現在は乗っていませんでした…。

pytorch-lightningでログを可視化したい【機械学習】

 

もとのスクリプト

 

デフォルトではTensorboardLoggerとのことで、公式Githubにはこのようなスクリプトとなっていました。

実行は以下です。

こちらをwandbで残せるよう変えていきます。

 

WandbLogger の使い方

pytorch lightningではカスタムロガーとしてWandbLoggerが用意されています。

使い方としては最低限以下を押さえておけばよさそうです。(他にもできることありそうですが、私はいったん必要なく試していません)

①loggerの定義

project引数で、今回の実験のプロジェクト名を指定します。

他にも、試行ごとにランダムで名前が付けられるものに関して、name引数で指定できたりします。

 

②ハイパーパラメータの記録

辞書型で、ハイパーパラメータを記述しておき、log_hyperparamsの引数に与えてあげることで記録できます。

 

③メトリクスの記録

※例として、学習データのロスを残す場合

メトリクスの記録に関しては、Class内で記述を行います。

logger.log_metrics()内に、キーバリューペアの形で渡してあげます。

 

④Trainerのインスタンス化のさいにLoggerを指定

logger引数に①で定義したWandbLoggerを指定してあげればOKです。

 

修正後スクリプト

というわけで、以下が修正後スクリプトです。

もとのスクリプトで不要になった部分(Tensorboardに残す記述)に関してはコメントアウトしています。

※メトリクスの記録をClassに追加。 #追加 の部分です。

 

Wandbのサイトから結果を確認

test_projectが作成されているはずですので、中を見に行くと無事に以下のように記録が残りました!

 

というわけで、PyTorch Lightningのログをwandbで残す方法でした。

どなたかのお役に立てば幸いです。

最後まで読んでいただきありがとうございました。