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

TensorboardXとは？

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

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

TensorboardXのインストール

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

pip install tensorboardX

基本的な使い方

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

from tensorboardX import SummaryWriter

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

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

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

記録できるデータの種類

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

数値データ

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

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

モデル構造

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

import torch.nn as nn
import torch

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

画像

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

import numpy as np

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

ヒストグラム

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

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

テキスト

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

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

TensorBoardでの可視化

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

tensorboard --logdir=runs

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

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

 

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

数値データ

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

モデル構造

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

画像

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

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

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

テキスト

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

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()

biasの可視化

weightの可視化

loss

CNNのトレーニング可視化

必要なセットアップ

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

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

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

ここでは、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)

モデル

シンプルな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

画像を編集する関数

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)

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

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

# モデル、損失関数、最適化手法の設定
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()

TensorBoardの起動と確認

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

tensorboard --logdir=runs

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

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

loss

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