PyTorchを利用したTensorBoardXの使い方について

1. TensorboardXとは？

TensorboardXは、PyTorchなどで機械学習モデルを構築する際に、トレーニングプロセスを可視化するためのライブラリです。今回はPyTorchと組み合わせて、データの可視化を行っていきます。

今回使用したtensorboardxのバージョンは

2. TensorboardXのインストール

以下のコマンドでTensorboardXをインストールします。

pip install tensorboardX

2.1. 基本的な使い方

まず、SummaryWriterを使ってログを記録します。

from tensorboardX import SummaryWriter

# ログ出力先を指定
writer = SummaryWriter(logdir='runs/experiment1')

# 必要なデータを記録
writer.add_scalar('Loss/train', 0.1, 1)  # スカラー値
writer.close()  # 記録終了

上記コードでは、runs/experiment1ディレクトリにログが保存されます。

3. 記録できるデータの種類

TensorboardXを使うと、以下のようなデータを記録できます。

3.1. 数値データ

損失や精度などの数値データを記録します。

for step in range(100):
    loss = 1 / (step + 1)  # 仮の損失値
    writer.add_scalar('Loss/train', loss, step)

3.2. モデル構造

モデルのネットワーク構造を可視化できます。

import torch.nn as nn
import torch

model = nn.Linear(10, 5)
dummy_input = torch.randn(1, 10)
writer.add_graph(model, (dummy_input,))

3.3. 画像

画像データをTensorBoardに記録して確認できます。

import numpy as np

image = np.random.rand(3, 100, 100)  # RGB画像
writer.add_image('Random Image', image, 0)

3.4. ヒストグラム

モデルのパラメータの分布を確認できます。

for step, param in enumerate(model.parameters()):
    writer.add_histogram('Model Param', param, step)

3.5. テキスト

簡単なテキストデータを記録します。

writer.add_text('Info', 'This is a sample text', 0)

4. TensorBoardでの可視化

データを記録した後、以下のコマンドでTensorBoardを起動します。

tensorboard --logdir=runs

ブラウザでhttp://localhost:6006にアクセスすると、データの可視化結果を確認できます。

データの項目は次の部分をクリックすることで選択できます。

 

先ほどのコードを実行したときには、次のようにデータを確認できます。

4.1. 数値データ

損失や精度などの数値データを記録します。

4.2. モデル構造

モデルのネットワーク構造を可視化できます。

4.3. 画像

画像データをTensorBoardに記録して確認できます。

4.4. モデルのパラメータの分布

モデルのパラメータの分布を確認できます。

4.5. テキスト

簡単なテキストデータを記録します。

5. PyTorchでのトレーニング可視化

以下は、PyTorchを使った基本的な線形モデルのトレーニングプロセスの例です。

import torch
import torch.nn as nn
import torch.optim as optim
from tensorboardX import SummaryWriter

# データ
x = torch.randn(100, 10)
y = torch.randn(100, 5)

# モデルと損失関数
model = nn.Linear(10, 5)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# SummaryWriterの初期化
writer = SummaryWriter(logdir='runs/experiment1')

# トレーニング
for epoch in range(100):
    model.train()
    optimizer.zero_grad()
    outputs = model(x)
    loss = criterion(outputs, y)
    loss.backward()
    optimizer.step()

    # スカラー値を記録
    writer.add_scalar('Loss/train', loss.item(), epoch)

    # モデルパラメータを記録
    for name, param in model.named_parameters():
        writer.add_histogram(f'{name}', param, epoch)

writer.close()

5.1. biasの可視化

5.2. weightの可視化

5.3. loss

6. CNNのトレーニング可視化

6.1. 必要なセットアップ

まず、ライブラリをインポートします。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import torchvision.transforms as T
from tensorboardX import SummaryWriter
import matplotlib.pyplot as plt

6.2. データセットとモデルの準備

ここでは、MNISTデータセットを使用します。

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)

6.3. モデル

シンプルなCNNモデルを使用します。

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(32 * 14 * 14, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = x.view(-1, 32 * 14 * 14)
        x = self.fc1(x)
        return x

6.4. 画像を編集する関数

matplotlibを利用して画像を作成して、ログに保存するための関数を作成します。

# Matplotlibで画像を描画し、TensorBoardに保存する関数
def log_images_with_labels(images, true_labels, predicted_labels, writer, epoch, num_images):
    # サンプル数を指定
    images = images[:num_images]
    true_labels = true_labels[:num_images]
    predicted_labels = predicted_labels[:num_images]

    # 描画用に列数を指定
    cols = 4
    rows = (num_images + cols - 1) // cols  # 必要な行数を計算
    fig, axes = plt.subplots(rows, cols, figsize=(12, 3 * rows))
    axes = axes.flatten()  # フラットなリストに変換

    for i in range(num_images):
        # TensorをPIL画像に変換
        image = T.ToPILImage()(images[i].cpu())

        # 画像を描画
        axes[i].imshow(image, cmap='gray')
        axes[i].axis('off')  # 軸を非表示

        # ラベルを描画
        true_label = true_labels[i].item() if isinstance(true_labels[i], torch.Tensor) else true_labels[i]
        pred_label = predicted_labels[i].item() if isinstance(predicted_labels[i], torch.Tensor) else predicted_labels[i]
        axes[i].set_title(f"True: {true_label}, Pred: {pred_label}", fontsize=10)

    # 空のプロットを隠す
    for i in range(num_images, len(axes)):
        axes[i].axis('off')

    plt.tight_layout()

    # TensorBoardに保存
    writer.add_figure(f'Epoch_{epoch+1}/Images', fig, global_step=epoch)

    # Matplotlibのメモリを解放
    plt.close(fig)

6.5. 学習ごとの画像データ記録

トレーニング中に学習ごとに画像を記録します。

# モデル、損失関数、最適化手法の設定
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# SummaryWriterの初期化
writer = SummaryWriter(logdir='runs/cnn_image_logging')

# トレーニングループ
for epoch in range(5):  # 5エポックのトレーニング
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, target)
        loss.backward()
        optimizer.step()

        # 損失の記録
        if batch_idx % 100 == 0:
            print(f'Epoch [{epoch+1}/5], Step [{batch_idx+1}/{len(train_loader)}], Loss: {loss.item():.4f}')
            writer.add_scalar('Loss/train', loss.item(), epoch * len(train_loader) + batch_idx)

    # 各エポックの終了時に画像を記録
    sample_images = data[:8]  # 最初の8枚
    true_labels = target[:8]
    _, predicted = torch.max(outputs[:8], 1)

    # ラベルを追加した画像を取得

    log_images_with_labels(sample_images, true_labels, predicted, writer, epoch,8)


    # テキストで予測結果を記録
    text_log = [f'True: {true_labels[i].item()}, Pred: {predicted[i].item()}' for i in range(len(true_labels))]
    writer.add_text(f'Epoch_{epoch+1}/Predictions', '\n'.join(text_log), epoch)

writer.close()

6.6. TensorBoardの起動と確認

以下のコマンドでTensorBoardを起動します。

tensorboard --logdir=runs

ブラウザでhttp://localhost:6006にアクセスすると、以下を確認できます。

• 各エポックごとの画像。
• 予測ラベルと真のラベルのテキスト記録。
• トレーニング損失の変化。

6.7. loss

6.8. ラベルと同時に表示された画像