kumpei.ikuta.me

深層学習のための実験管理スイート Weights & Biases

最近は深層学習の実験を Weights & Biases （略称：WandB）で管理しており、とても便利なので紹介します。

わかる人向けに説明すると、Tensorboard ライクなロギングをしつつ、実験（パラメータ + ログ + 最終結果）を管理するツール、といった感じです。

実験は Run という単位で管理され、各 Run は名前や ID 、タグなどに加え、任意のカラム（例：accuracy_test）を持つことができます。そして、これらのカラムの値に応じて、Run を絞りこんだりソートしたりできます。

Run一覧ページのスクショ
Run一覧ページのスクショ

MLFlow と何が違うの？

MLFlow は MLFlow 式のパラダイムがあって、全てそこに乗っかっていく、みたいな設計になっています。一方 WandB はどちらかというと既存のコードに最小限の変更を加えるだけで使えることを意識しているようです。ソフトウェア設計としては WandB のほうが優れていますが、MLFlow を使ったコードは誰が書いても似たような見た目になるというメリットがあります。チーム開発をしていて、かつ既存のコードベースが無いのであれば MLFlow も有効な選択肢だと思います。

つかいかた

まずはアカウントを作りましょう

PyPI からクライアントをインストールして、ログイン。アクセストークンは User Settings で見れます。

$ pip install wandb
$ wandb login アクセストークン

ここからは Python でゴリゴリとコーディング……の前に、絶対に知っておくべきコマンドを紹介します。wandb on と wandb off です。

$ wabdb on
$ wandb off

読んで字のごとく WandB を一時的にオンオフするコマンドで、たとえば「ちょっとここのテンソルの形状を確認したいな」といったときに、いちいち Run が作られてデータが送られて……となるのを防ぐことができます。

さて、それでは Python を書いていきます。まずはコード内で init して、プロジェクト（Run をまとめるもの）を指定します。またこのとき whateveriwant というプロジェクトを Web 上で作っておく必要はないです。なければ勝手に作られます。

wandb.init(project='whateveriwant')

次にハイパーパラメータを保存しときましょう。これは必須ではないですが、良い結果が出たときに、どういうパラメータだったっけ？というのを参照するのに役立ちます。再現性は大事。

wandb.config.learning_rate = 1e-4
wandb.config.job = None

スカラ値で表現できないハイパーパラメータ（学習のアルゴリズムやネットワーク構造など）については、Run の .notes にテキストで書いておくと良いでしょう。ここにはマークダウンも書けます。

wandb.run.notes = "SRResNet + Deconv"

また、PyTorch の場合はネットワークを渡してあげると勾配のログを自動で取ってくれます。

wandb.watch(model)

ここまでやったら、あとはロギングするだけです。

for epoch in range(1, num_epochs + 1):
    wandb.log({
        'loss/train': 0,
        'loss/val': 0,
        'acc/train': 1,
        'acc/val': 1
    })

wandb.Image に torch.Tensor やそのリストを渡すと画像も表示できます。ただし Pillow が必要（Pillow-SIMD じゃダメらしい）。

wandb.log({
    'selfie': wandb.Image(torch.rand(64, 64)),
    'friend_photos': [wandb.Image(torch.zeros(64, 64)) for _ in range(8)]
})

HTML も出せます。これが地味に便利で、matplotlib で描いた複雑な図を HTML にエクスポートして WandB 上で見る、とかができます。嬉しい。

wandb.log({
    'myhomepage': wandb.Html('<em>Hello</em>')
})

リファレンス によれば、点群データや分子構造など、ここで紹介していない様々な物をロギングできるようです。いつか使ってみたいですね。

以上、実験管理スイート Weights & Biases の紹介でした。

back to index