PyTorchのTensorBoardサポートを試してみる #Python

はじめに

PyTorchのv1.1.0からオフィシャルのTensorBoardサポート機能が追加されました。torch.utils.tensorboard にあるSummaryWriter を使うことで、PyTorch を使っているときでも、学習ログなどの確認にTensorBoard を活用することができます。

この記事では、このSummaryWriter の使い方を簡単に紹介したいと思います。

この記事について (20/03/02)

この記事を最初に投稿したのはPytorch v1.1.0 が出てすぐの頃(19/5/3)で、情報が古くなっている部分がちょこちょこあったため内容をアップデートしました。ついでに記事の構成も少し変更しています。

当初は使い方に関する情報が少なかったのですが、現在ではPytorch 公式ホームページに説明がまとまっています（下のリンク二つ）。この記事は流し読みぐらいにして、リンク先の説明をしっかり読むといいと思います笑。

準備編

PyTorch とTensorBoard をインストールしましょう。
（既にどちらもインストールされていて、かつバージョンが古すぎなければここは読み飛ばして大丈夫です。）

どちらもpip (Anaconda ユーザーならconda) でインストールできます。

（Anaconda ユーザへの注意: よく言われているようにpip とconda を混ぜて使うと思わぬトラブルを引き起こすことがあるので、Anaconda を使っているならconda を使うのが無難です。）

Pytorch のインストール

公式ページでの説明に従ってください。

環境によってインストールコマンドが違います。
TensorBoard サポートはv1.1.0 以降からです。特に理由がなければ最新のものを使えばよいと思います。

TensorBoard のインストール

公式での説明が見当たらなかったのですが、

pip install tensorboard (or conda install tensorboard)

でとりあえず入ります。
ただ、TensorBoard のgithub を見るとTensorFlow が入っていないと機能が制限されると書いてあるのでTensorFlow をインストールした方がいいのかもしれません。（TensorFlow をインストールするとTensorBoard は依存で勝手についてきます。）

TensorFlow のインストールは公式での説明をご参照ください。
TensorBoard のバージョンがあまり古いとだめかもしれません。
- 参考までに、PyTorch v1.1.0 のときはTensorBoard はv1.14 以上が要求されていました。

SummaryWriter の使い方

さて、本題のSummaryWriterの使い方についてです。
実用的な例（ニューラルネットの学習ログをとるなど）を持ち出すと関係ない部分が長くなってしまうので、ここでは単純化したミニマルな例をベースに説明したいと思います。

ミニマルな使用例

適当に作ったnumpy配列をTensorBoard でプロットしてみます。

basic_usage1.py

import numpy as np
from torch.utils.tensorboard import SummaryWriter


# データを作る
np.random.seed(123)
x = np.random.randn(100)
y = x.cumsum()  # xの累積和

# ここではmatplotlib での以下に相当するものをTensorBoard で表示します。
# t = np.arange(100)
# plt.plot(t, x)
# plt.plot(t, y)

# log_dirでlogのディレクトリを指定
writer = SummaryWriter(log_dir="./logs")

# xとyの値を記録していく
for i in range(100):
    writer.add_scalar("x", x[i], i)
    writer.add_scalar("y", y[i], i)

# writerを閉じる
writer.close()

上のコードを実行すると、カレントディレクトリに./logs が作られます。
ターミナル（or コマンドプロンプトなどなど）で

tensorboard --logdir [logs へのパス]

を実行しましょう。デフォルトの設定のままだとlocalhost:6006 にアクセスすることで次のようなプロットを見ることができます：

プロットを拡大したり、スムージングをかけたりできてとても良いです。
ちなみに、この例では実行時間が短すぎてわかりませんが、プログラム実行中でもlocalhost:6006 にアクセスすることでプロットが更新されていくのをモニタリングできます。
また、生データをCSV やJSON 形式でダウンロードすることもできます。

解説

さて上のコードに説明を加えていきましょう。
やっていることは単純で、

writer = SummaryWriter(log_dir="./logs") でSummaryWriter のインスタンスを作る。
- ここでlog_dir にTensorBoard 用のファイルを吐き出すディレクトリを指定
writer.add_scalar で値を記録
最後にwriter.close() で閉じる

という流れです。

add_scalar メソッドの引数はadd_scalar(tag, scalar_value, global_step=None, walltime=None)のようになっています:

tag
- データの名前
- 後述するようにtag の付け方でプロットをグルーピングできます
scalar_value
- 記録したい値
global_step
- 何番目の値かを指定
- 現在のepoch数など

(walltimeはどういう時にいじるのだろう? )

その他のメソッドについて

ここではadd_scalar メソッドを取り上げましたが、add_something という形式のメソッドが他にもたくさんあります。例えばadd_image とかadd_histogram とか。

どれもだいたいadd_something(tag, data, global_step) という感じで使えるので、気になった人はDocs を見てみましょう。

Tips

上ではとても単純な例しか取り上げてませんが、これでもうSummaryWriter を使えるようになったと思います！
ただ、ここまでで記事を終えてしまうのも味気ないので、ここからは知っておくと便利なtips を紹介していきます。

tag によるプロットのグルーピング

まずは予告しておいたtag によるプロットのグルーピングについてです。

上の例ではtag に安直な名前を指定していました。
しかし、ログをとる対象が増えてくると、上のように全てが別枠で表示されてしまうと見にくいです。

例えば、ニューラルネットの学習を記録する時なら、loss やaccuracy ... etc. といった評価指標、しかもそれらのtrain, validation それぞれでの値と監視対象が多いですよね。

実は、tag の名前の付け方でプロットをグルーピングできます。
具体的には、tag="group_name/value_name" のようにすることで、group_name ごとにプロットの枠を分けてくれます。

今回もとっても単純化した例で試してみましょう。

basic_usage2.py

import numpy as np
from torch.utils.tensorboard import SummaryWriter


# ログをとる対象を増やしてみる
np.random.seed(111)
x1 = np.random.randn(100)
y1 = x1.cumsum()

x2 = np.random.randn(100)
y2 = x2.cumsum()

writer = SummaryWriter(log_dir="./logs")

# tag = "group_name/value_name" の形式に
for i in range(100):
    writer.add_scalar("X/x1", x1[i], i)
    writer.add_scalar("Y/y1", y1[i], i)
    writer.add_scalar("X/x2", x2[i], i)
    writer.add_scalar("Y/y2", y2[i], i)

writer.close()

TensorBoardで見てみるとこんな感じになります。

ログディレクトリの構成

ハイパーパラメータ調整などでたくさん実験を行うときは、ログの親ディレクトリを決めておいて、各実験ではそのサブディレクトリをlog_dir に指定するようにするとTensorBoard でまとめて確認しやすくなります。

例えばこんな感じ:

logs       # ログの親ディレクトリ
├── exp1　 # 実験1のログ（内容に応じた名前を付けると後で見やすいです）
├── exp2
└── exp3

実験1のログをSummaryWriterで書き出す時は、

writer = SummaryWriter(log_dir="logsへのパス/exp1")

とすれば上のようなディレクトリ構成になりますね。

このようなディレクトリ構成にしておいて、TensorBoard を起動する際に親ディレクトリを指定するとまとめてプロットを表示してくれます。

tensorboard --logdir="logsへのパス"

例

optimizer を色々変えてみて、結果を比べてみた例です。
ログディレクトリの構成は次のようになっています:

logs
├── adagrad   # どのoptimizerでの実験か分かるような名前にしました
├── adam
├── amsgrad
├── sgd_lr0.01
├── sgd_lr0.1
├── sgd_lr0.1_momentum0.9
├── sgd_lr1.0
└── sgd_lr1.0_momentum0.9

全部まとめて表示してくれて比較しやすいですね。

tensorboardX との関係

PyTorch の中にあるSummaryWriter とほぼ同じものを提供してくれるパッケージとしてtensorboardX があります。ちなみに、名前の由来は「tensorboard for X」らしいです。

tensorboardX にもSummaryWriter というクラスがあって、使い方はPyTorch にあるものとほぼ一緒です。古いバージョンのPyTorch を使いたい時などはtensorboardX を代わりに使うといいと思います。

PyTorchのgithubレポジトリを覗いてみると、tensorboardXの作者によるプルリクTensorBoard support within PyTorch #16196から、PyTorchのTensorBoardサポートにつながっていったみたいですね。

まとめ

以上、PyTorchのSummaryWriterの簡単な紹介でした。
プロットって意外と時間をとられがちなので、こういうところが便利になっていくのはうれしいですね。