TensorBoardのGoogle Colabでの設定と、PyTorchでの使い方

PyTorch 1.1より、TensorBoardによる監視機能がデフォルトで入った(ただし、1.1.0ではExperimentalで、1.2.0より正式版)。ここでは、Google ColabでのTensorBoardの設定方法および、PyTorchでの使い方について説明する。

Google ColabでのTensorBoardの設定方法

TensorBoardは、学習状況の監視ツールである。Google Colabでは、2つの方法で使うことが出来る。一つは、Colabのセルで表示する方法であり、もう一つは、ColabにURLで接続する方法である。
まず、PyTorch 1.1用のTensorBoardに更新する。(TensorBoard 1.14以降が必要であるため、いずれかのコマンドが必要。)

!pip install tensorboard==1.14.0

!pip install tb-nightly

この後、セル方式とURL方式に手順が分かれる。

Google Colabのセルで表示する方法

以下のようにしてTensorBoardを起動する。なおlogdirをrunsとしているのは、PyTorchのSummaryWriterのデフォルトのためである。TensorBoardColabと共有できるようにするためには、logdirをGraphにする必要がある。

%load_ext tensorboard
%tensorboard --logdir ./runs

なお、%tensorboardのGoogle Colab/Jupyter notebookの組み込みコードは下に記載している。(ちょっと上級者向け)

ngrokを用いて、外部URLでTensorBoardを動かす方法

以下の手順により、URLを介してTensorBoardで表示することが出来る。下記では、tensorboardを動かしており、port 6006を使っていることがわかる。

LOG_DIR = './runs'
get_ipython().system_raw(
    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 &'
    .format(LOG_DIR)
)
!wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip
!unzip ngrok-stable-linux-amd64.zip
get_ipython().system_raw('./ngrok http 6006 &')
!curl -s http://localhost:4040/api/tunnels | python3 -c \
    "import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])"

PyTorchからの監視(TensorBoardへのデータ書込み)

torch.utils.tensorboard

以下のコードを実行すると、image, graph, textおよびembedding (TensorBoardではPROJECTOR) が表示される。なお、各関数の使い方は、torch.utils.tensorboardでは、サンプルコードが少ないので、tensorboardXを見たほうが良いかもしれない。

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms
import keyword

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

# Scalar
import numpy as np
for n_iter in range(10):
    writer.add_scalars('data/scalar_group', {"xsinx": n_iter * np.sin(n_iter),
                                                         "xcosx": n_iter * np.cos(n_iter),
                                                         "arctanx": np.arctan(n_iter)}, n_iter)

# Image (図)
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
images, labels = next(iter(trainloader))
grid = torchvision.utils.make_grid(images)
writer.add_image('images', grid, 0)

# Video (動画 IMAGEタブに表示される)
n_iter =0
writer.add_video('myVideo', torch.rand(16, 48, 1, 28, 28), n_iter)

# Graph　(モデル構成図)
model = torchvision.models.resnet50(False)
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
writer.add_graph(model, images)

# Audio (音声)
sample_rate = 44100
n_iter = 0
x = torch.zeros(sample_rate * 2)
writer.add_audio('myAudio', x, n_iter)

# Text
writer.add_text('lstm', 'This is an lstm', 0)

# Embedding (単語等のベクトル表現)
meta = []
while len(meta)<100:
    meta = meta+keyword.kwlist # get some strings
meta = meta[:100] 
for i, v in enumerate(meta):
    meta[i] = v+str(i)
label_img = torch.rand(100, 3, 10, 32)
for i in range(100):
    label_img[i]*=i/100.0
writer.add_embedding(torch.randn(100, 5), metadata=meta, label_img=label_img)
writer.close()

さまざまな測定を見る。

PyTorchに含まれているtest_tensorboard.pyを用いて、どのような監視が出来るかをみることが出来る。以下に書込みに必要なコマンドを示す。以下のスクリプトで、sedを使っているのは、デフォルトで出力するrunsディレクトリをテスト直後に消してしまうことを防ぎTensorBoardで見るために追加している。

!git clone http://github.com/pytorch/pytorch
!pip install ninja
!pip install matplotlib
!pip install tensorflow
!pip install torchvision==0.3.0
!pip install tb-nightly
!sed -i 's/runs/sruns/g' pytorch/test/test_tensorboard.py
!python3 pytorch/test/run_test.py -f tensorboard  -v
!rm -r runs
!cp -r pytorch/test/runs .

監視機能比較

監視機能に関しては、torch.utils.tensorboardが圧倒的に高機能である。

機能	torch.utils.tensorboard	TensorFlow/TensorBoard	TensorBoardColab
scalar(スカラー値)	○	○	○
histogram(ヒストグラム)	○	○	x
distribution(ヒストグラム)	○	○	x
image(図 IMAGE TAB)	○	○	○
video(動画 IMAGE TAB)	○	○	×
graph(学習モデル GRAPH TAB)	○	○	x
audio(音声 AUDIO TAB)	○	○	x
text(テキスト TEST TAB)	○	○	x
embedding(ベクトル図即ち分散表現 PROJECTOR TAB)	○	○	x

その他　(本記事を書いてからの変更点)

PyTorch 1.1.0リリース後の修正

• Exception when adding custom scalars in tensorboard #20579
  • 2019/5/24修正 add_custom_scalars(layout)が適切に動かない件
  • なお、本関数のテストセットは見かけなかった。参考までにMergeFromは、protocol bufferのメッセージ

TensorBoard 1.14.0のリリース

2019年6月15日リリース

• TensorBoard 1.14.0
• notebook: relocate to %load_ext tensorboard #2119
  • 以前は、%load_ext tensorboard.notebookが必要だったが、%load_ext tensorboardに名前空間が変更になった。

その他

Tensorboardを使うtorch.utils.tensorboard以外のモジュール

Tensorboardでのモデルのグラフ表示は出来るが、それ以外に、Caffe2でのグラフ表示や、torch.jitのグラフ表示ツールが公開されている。ただし、tensorflowのGraphExecutorの取り込みが前提なので、デバッグ用とかの特殊用途である。(PyTorchでの通常利用では、tensorflowは使わない。)

• pytorch/caffe2/contrib/tensorboard/
• pytorch/torch/contrib/

参考資料

• (torch)TORCH.UTILS.TENSORBOARD
  • Using TensorBoard with PyTorch 1.1.0
  • tensorboardX
    • torch.utils.tensorboardの前身
• ngrok
• tensorboard/releases
• tensorflowcolab
  • How to use Tensorboard with PyTorch in Google Colab
  • 初期化パラメータが記載されているコード
  • ngrok

メモ

TensorBoardのプラグイン

%load_ext tensorboardで、組みこんでいる箇所を簡単に記載する。以下の処理の延長で、_start_magic関数が呼び出されて(Cell版の)TensorBoardが稼働する。
現状では、稼働する操作%load_extは使えるが、止める操作%unload_extは実装されていない。

• def load_ipython_extension(ipython)
  • def _load_ipython_extension(ipython)
• IPython extensions
• TensorBoardのJupyter/Colabプラグインコード
  • notebook: add Jupyter integration (#1813)
  • notebook: add Colab integration shell (#1822)

TensorBoardColabに関するメモ

PyTorchでもtensorboardcolabが使えるが、推奨ではないので文末に記載する。

環境構築方法

TensorBoardColabは、Google Colabのインスタンスにデフォルトで入っている。これにより、URLで接続できる。機能は以下の通りである。

• ngrokによるURLポート公開
• tensorboardサーバ起動
• データの書き込み機能(/content/Graph)

初期化は、以下の通り行う。これにより、ngrokとtensorboardが起動する。なお、このスクリプトを実行すると、接続すべきURLが表示される。

from tensorboardcolab import TensorBoardColab
tbc = TensorBoardColab()

なお、TensorBoardColabの詳細設定は、ソースコードに記載されている。(README.mdは見かけない。)graph_path=./Graphとなっているので、それを書き換えて設定を変更することが出来る。

データ書込み

TensorBoardColabの初期化は、すでに記載した。ここではデータ書込みを説明する。
データ書込みは、以下の機能が提供されている。

tbc.save_image(title, image)
tbc.save_value(graph_name, line_name, epoch, value)
tbc.flush_line(line_name)
tbc.close()