Pytorch で TensorBoard を使用する

「TensorBoard」は、TensorFlow に付属した可視化を助けるツールです。

ですが、「Pytorch」でも、公式にサポートされていますので、Pytorch を使用すれば、TensorFlow と同じく、TensorBoard でグラフの作成などを行うことができます。

今回は、この Pytorch を使用して「Google Colaboratory」の Notebook 上で TensorBoard を起動し、活用する方法について解説したいと思います！

TensorBoard を使用するために TensorBoardX をセットアップ

1!pip install tensorboardX

TensorBoard を使用する流れ

TensorBoard を使用する流れは、以下の通りです。

1. tensorboardX.SummaryWriter() を使用し、ログ保存用のファイルディレクトリを作成する
2. add_something(tag name, object, iteration number) を用いて、各データに応じたログの保存を行う
3. SummaryWriter を閉じる
4. TensorBoard を起動し、ログデータを表示する

1from tensorboardX import SummaryWriter
2
3#ログディレクトリを指定しない場合は './runs/(CURRENT_DATETIME_HOSTNAME)' が自動で生成されます
4writer1 = SummaryWriter()
5
6#または任意のログディレクトリ（logdir）を指定できます。以下は'./logs/image'が生成されます。
7writer2 = SummaryWriter(logdir='logs/image')
8
9#writer1の様にlogdirを指定していない場合、commentを指定するとデフォルトのログディレクトリの後に続いてcommentが追加されます。
10writer3 = SummaryWriter(comment='loss function')

SummaryWriter() に任意のログディレクトリを指定でき、指定しない場合は、デフォルトのディレクトリが作成されます。

ログを保存する場所を整理したり、複数のデータや設定を比較したりするときなどに使い分けることが可能です。

定義したログディレクトリ

ディレクトリを定義したら、次に、ログの保存を行っていきます。

ログ保存のフォーマット

ログを保存するフォーマットは、add_something(tag name, object, iteration number) となっています。

そして、something の部分は、データの種類に応じて、

1. add_scalar
2. add_image, add_images
3. add_histogram
4. add_figure
5. add_graph
6. add_embedding
7. add_pr_curve

add_image, add_images で画像を表示する

MNIST の画像データを TensorBoard で表示してみましょう。

MNIST の画像データセットを用意

1#データの用意（MNIST）
2
3import torch
4import torchvision
5from torchvision import datasets, transforms
6
7# transforms
8transform = transforms.Compose([transforms.ToTensor(), 
9                                transforms.Normalize((0.5,), (0.5,))])
10
11# dataset
12train_dataset = datasets.MNIST(root='./content/',
13                               train=True,
14                               transform=transform,
15                               download=True)
16
17test_dataset = datasets.MNIST(root='./content/',
18                              train=False,
19                              transform=transform)
20
21
22# Data Loader
23batch_size = 1000
24train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
25                                           batch_size=batch_size,
26                                           shuffle=True)
27
28test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
29                                          batch_size=batch_size,
30                                          shuffle=False)

make_grid で画像データを作成し、writer に保存する

1images, labels = next(iter(train_loader))
2
3#5x5で画像を並べます
4img_grid = torchvision.utils.make_grid(images[:25], nrow=5)
5
6#'logs/image' にimg_gridデータを'mnist_images'というタグ名で保存する
7writer = SummaryWriter(logdir='logs/image')
8writer.add_image('mnist_images', img_grid)
9writer.close()

make_grid で作成した画像データを add_image() で「writer」に保存します。

「writer」への書き込み後は、 close()  で SummaryWriter を閉じます。

TensorBoard を起動し、img_grid を表示する

では、TensorBoard を起動し、保存した img_grid を表示してみます。

Colaboratory では、以下のコードで「TensorBoard notebook」の読み込みと起動を行います。

1#TensorBoard notebook の読み込み（一度でOK）
2%load_ext tensorboard
3
4#TensorBoard起動（表示したいログディレクトリを指定）
5%tensorboard --logdir=logs

実行すると、Notebook 内で TensorBoard が起動し、保存した img_grid の画像が指定したタグ名で表示されています。

img_grid の画像

add_image() には、タグ名（tag）と画像データを渡します。

画像データは、uint8(0, 255）、もしくは、float32（0, 1）の Tensor です。

形状は、[channel, height, width] がデフォルトですが、dataformats に 'HW' や 'HWC' など渡すことで対応できます。

この img_grid は、torchvision の make_grid で画像を並べたものを add_image() で表示しています。

add_images() で複数枚の画像を並べて表示する

ですが、add_images() を使えば、複数枚の画像 [n, channel, height, width] を1枚ずつ並べることができます。

1#16枚の画像を並べる
2img_batch = images[0:16]
3
4#add_imagesを使用し、別のタグ名で保存
5writer.add_images('mnist_images_2', img_batch)
6writer.close()

同じ様に、add_images でログを保存したら、先ほど起動した TensorBoard のリロードアイコンで更新すると新しく追加したログデータを表示できます。

並べて表示した16枚の画像

これらは、画像を表示しただけですが、add_embedding を使用することで、より TensorBoard の威力を発揮することができます。

add_embedding による次元削減と可視化

データ分析では、高次元データの次元削減や特徴空間上への可視化を、以下などの手法を使って行います。

1. PCA（Principal Component Analysis）
2. t-SNE（t-distributed Stochastic Neighbor Embedding）
3. UMAP（Uniform Manifold Approximation and Projection）

add_embedding() では、これらを TensorBoard 上で簡単に行うことができます。

add_embedding() に渡すのは、プロットする [n : プロット数, d : 次元数] の Tensor、リストのラベルデータ、ラベルの画像データ[N,C,H,W] です。

今回は、MNIST の784次元、1000枚の画像をプロットしています。

1mat_img = images.view(-1, 28*28)
2meta_labels = [str(x) for x in labels.numpy().tolist()]
3
4#add_embeddingでログを保存（with文を使用）
5with SummaryWriter(logdir='logs/projector') as w:
6  w.add_embedding(mat_img, metadata=meta_labels, label_img=images)

【PROJECTOR】タブが作成され、3次元空間上に MNIST の手書き数字画像が描画されます。

PCA を用いて可視化

このプロジェクター上では、自由に軸を動かしたり、ズームしたりできます。

また、画像をクリックすると近い特徴量をもつ周りの画像を表示してくれたりします。

「PCA」の他に、「t-SNE（下画像）」や「UMAP」タブがあり、それぞれの方法で「学習」、「クラスタリング」、「可視化」を行うことができます。

そのほか、見た目などを設定することができるので、試してみてください。

t-SNE を用いて可視化

add_graph でネットワークモデルをグラフ化する

次に、学習ネットワークモデルを構築し、それを TensorBoard でグラフ化したいと思います。

学習ネットワークモデルを構築

1from torch import nn, optim
2import torch.nn.functional as F
3
4class Net(nn.Module):
5  def __init__(self):
6    super(Net, self).__init__()
7    self.conv1 = nn.Conv2d(1, 6, 5)
8    self.pool = nn.MaxPool2d(2, 2)
9    self.conv2 = nn.Conv2d(6, 16, 5)
10    self.fc1 = nn.Linear(16*4*4, 120)
11    self.fc2 = nn.Linear(120, 84)
12    self.fc3 = nn.Linear(84, 10)
13  
14  def forward(self, x):
15    x = self.pool(F.relu(self.conv1(x)))
16    x = self.pool(F.relu(self.conv2(x)))
17    x = x.view(-1, 16*4*4)
18    x = F.relu(self.fc1(x))
19    x = F.relu(self.fc2(x))
20    return self.fc3(x)
21    
22
23model = Net()
24criterion = nn.CrossEntropyLoss()
25optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

構築したネットワークモデルをグラフ化

ここで、add_graph() を使用するのですが、TensorBoardX で add_graph を使用して作成すると、不完全なグラフになってしまったり、表示がされないという不具合がみられました。

ただ、通常の torch.utils.tensorboard の SummaryWriter クラスを使用して、同様に add_graph() で作成すれば問題なく Tensorboard で表示することができました。

1#TensorBoardXを使用した場合
2with SummaryWriter(log_dir='logs/model') as w:
3  w.add_graph(model, (images.requires_grad_(True),))
4
5#torch.utils.tensorboardを使用した場合
6import torch.utils.tensorboard
7with torch.utils.tensorboard.SummaryWriter(log_dir='logs/model_2') as w:
8  w.add_graph(model, images)

add_graph は、ネットワークモデルと学習データを渡すことで、そのネットワークがどういう構造で構成されているかを模式的に表示させることができます。

データの流れやネットワークの構造を視覚的に確認できるので、モデルを整理して見直すことができます。

グラフ化したネットワーク構造

add_scalar による学習の可視化

ここからは、モデルの学習を行い、add_scalar() を用いて損失関数値のグラフ化を行いたいと思います。

モデルの学習を実施

1writer = SummaryWriter(logdir='logs/train_loss')
2
3running_loss = 0.0
4for epoch in range(0,10):
5  model.train()
6  for i, (data, target) in enumerate(train_loader):
7    i += 1
8    optimizer.zero_grad()
9    output = model(data)
10    loss = criterion(output, target)
11    loss.backward()
12    optimizer.step()
13
14    running_loss += loss.item()
15    if i % 5 == 0:
16      print('Train Epoch: {} | Batch Status: {}/{} ({:.0f}%) | Loss: {:.6f}'.format(
17          epoch, i * len(data), len(train_loader.dataset),
18          100. * i / len(train_loader), loss.item()
19      ))
20      # 損失関数の値をログに保存
21      writer.add_scalar('training_loss',
22                        running_loss / 5,
23                        epoch * len(train_loader) + i)
24      
25      running_loss = 0.0
26
27writer.close()
28print('Finished Training')

学習の進行を可視化

add_scalar() は、スカラー値のシンプルなグラフ化を行うことができます。

今回は、学習中の損失関数の値を「5」イテレーション毎に、ログを保存する形でグラフ化をしました。

【SCALARS】というタブに、損失関数値のグラフ（縦軸が損失、横軸はイテレーション）が表示されています。

損失関数値のグラフ

このように、対数グラフでの表示もできます。

イテレーション数が「250」を超えると、損失関数の値が減少し、学習が進んでいる様子が分かりますね。

マウスをグラフに乗せると、ピンポイントの数値を確認できたり、スムージングやデータのダウンロードなどを行うことができます。

add_pr_curve による PR 曲線描画

最後に、学習済みモデルの検証のため、テストデータによる精度算出と add_pr_curve() を用いて TensorBoard 上に PR（Precision - Recall）曲線を描画します。

テストデータでモデルの精度を算出

1total = 0
2correct = 0
3class_probs = []
4class_labels = []
5classes = [str(x) for x in range(0, 10)]
6model.eval()
7with torch.no_grad():
8  for i, (data, target) in enumerate(test_loader):
9    output = model(data)
10    probs_batch = [F.softmax(x, dim=0) for x in output]
11    _, preds_batch = torch.max(output, 1)
12    total += target.size(0)
13    correct += (preds_batch ==  target).sum().item()
14    class_probs.append(probs_batch)
15    class_labels.append(target)
16
17  print("Accuracy : {:.2f} %".format(correct / total * 100))
18
19  test_probs = torch.cat([torch.stack(batch) for batch in class_probs])
20  test_labels = torch.cat(class_labels)
21
22  writer = SummaryWriter(log_dir='logs/pr_curve')
23  #各クラス毎にadd_pr_curve（）によりログデータを保存
24  for i in range(10):
25    test_labels_binary = test_labels == i
26    writer.add_pr_curve(classes[i], test_labels_binary, test_probs[:, i], 0)
27
28  writer.close()

算出したモデルの精度

1#出力
2Accuracy : 91.95%

TensorBoard 上に PR 曲線を描画

add_pr_curve() に与えるデータは、テストデータのバイナリーラベル（0、1、もしくは、True、False）、モデルの予測確率値（0~1）で、どちらも Tensor です。

これを、各クラス毎（今回は0〜9）にグラフ化していくと、それぞれの PR 曲線が得られます。

また、閾値（num_thresholds）は、デフォルトで「127」であり、自由に設定してグラフを確認することができます。

PR 曲線のグラフ

さいごに

TensorBoard は、便利であると同時に、視覚に訴えるので楽しみながら使うことができます。

データからネットワークモデルの学習・検証まで、一連の機械学習プロセスを場面に応じて可視化でき、より効率的で直感的な分析ができるようになります。

ぜひ、他の機能についても調べてみてください！

こちらの記事もオススメ！