Pytorchのログをtensorboardで可視化する

SummaryWrite()のデフォルトのログ保存先はrunsディレクトリ

デフォルトだと反映される時間が遅いので、flush_secs=1オプションをつける
writer = SummaryWriter(flush_secs=1)

以下で起動する

$ tensorboard --logdir runs

ソース

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
import numpy as np
from torch.utils.tensorboard import SummaryWriter

num_epochs = 10         # 学習を繰り返す回数
num_batch = 100         # 一度に処理する画像の枚数
learning_rate = 0.001   # 学習率
image_size = 28*28      # 画像の画素数(幅x高さ)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

transform = transforms.Compose([
    transforms.ToTensor()
    ])

train_dataset = datasets.MNIST(
    './data',               
    train = True,           
    download = True,        
    transform = transform   
    )

test_dataset = datasets.MNIST(
    './data', 
    train = False,
    transform = transform
    )

train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size = num_batch,
    shuffle = True)
test_dataloader = torch.utils.data.DataLoader(
    test_dataset,     
    batch_size = num_batch,
    shuffle = True)

class Net(nn.Module):
    def __init__(self, input_size, output_size):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_size, 100)
        self.fc2 = nn.Linear(100, output_size)

    def forward(self, x):
        x = self.fc1(x)
        x = torch.sigmoid(x)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

model = Net(image_size, 10).to(device)
criterion = nn.CrossEntropyLoss() 
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate) 
model.train()  
writer = SummaryWriter(flush_secs=1)
for epoch in range(num_epochs): # 学習を繰り返し行う
    loss_sum = 0
    correct = 0

    for inputs, labels in train_dataloader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        inputs = inputs.view(-1, image_size) # 画像データ部分を一次元へ並び変える
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss_sum += loss
        loss.backward()
        optimizer.step()
        pred = outputs.argmax(1)
        correct += pred.eq(labels.view_as(pred)).sum().item()

    print(f"Epoch: {epoch+1}/{num_epochs}, Loss: {loss_sum.item() / len(train_dataloader)}")
    writer.add_scalar("loss_sum", loss_sum, epoch)
    accuracy = 100*correct/len(train_dataset)
    writer.add_scalar("train accuracy", accuracy, epoch)
    torch.save(model.state_dict(), 'model_weights.pth')

model.eval()  # モデルを評価モードにする
loss_sum = 0
correct = 0
with torch.no_grad():
    for inputs, labels in test_dataloader:
        inputs = inputs.to(device)
        labels = labels.to(device)
        inputs = inputs.view(-1, image_size) # 画像データ部分を一次元へ並び変える
        outputs = model(inputs)
        loss_sum += criterion(outputs, labels)
        pred = outputs.argmax(1)
        correct += pred.eq(labels.view_as(pred)).sum().item()

print(f"Loss: {loss_sum.item() / len(test_dataloader)}, Accuracy: {100*correct/len(test_dataset)}% ({correct}/{len(test_dataset)})")