用tensorboard支持pytorch训练可视化

在工作用了tensorboard来可视化模型训练过程后，发现还挺香的。另外pytorch也正式支持tensorboard了，这里记录一下。

前置条件

安装tensorboard：

pip install tensorboard

实现步骤

1. 指定tensorboard输出日志：writer = SummaryWriter(log_dir=LOG_DIR)
2. 将模型和数据集添加到writer中：writer.add_graph(model, images.to(device))
3. 记录过程数据指标：writer.add_scalar('Test Loss', avg_loss, epoch)
4. 当模型开始训练后，启动tensorboard：tensorboard --logdir=runs。打开链接就能看到模型过程指标了：http://localhost:6006/

代码示例

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
from datetime import datetime

# 1. 设置参数
BATCH_SIZE = 64
EPOCHS = 100
LEARNING_RATE = 0.001
NUM_CLASSES = 10
LOG_DIR = "runs/fashion_mnist_experiment_" + datetime.now().strftime("%Y%m%d_%H%M%S")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 2. 准备数据集
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_set = torchvision.datasets.FashionMNIST(
    root='./data',
    train=True,
    download=True,
    transform=transform
)

test_set = torchvision