Pytorch学习记录-Pytorch可视化使用tensorboardX

Pytorch学习记录-Pytorch可视化使用tensorboardX

在很早很早以前（至少一个半月），我做过几节关于tensorboard的学习记录。

https://www.jianshu.com/p/23205a7921cd
https://www.jianshu.com/p/6235c1ecde67
https://www.jianshu.com/p/2b24454b0629
https://www.jianshu.com/p/0080047e5456

迟迟没有转到Pytorch的原因也是tensorflow的可视化做的好，不过现在Pytorch也支持了，在教程里有，学习一个。
在本教程中，使用简单的神经网络实现MNIST分类器，并使用TensorBoard可视化训练过程。在训练阶段，我们通过scalar_summary绘制损失和准确度函数，并通过image_summary可视化训练图像。此外，我们使用histogram_summary可视化神经网络参数的权重和梯度值。
pytorch使用tensorboard有三种方法：

昨天看了一下余霆嵩的教程，推荐使用tensorboardX，使用比logger更方便一些。
注意看注释就行了，这里我没有生成更复杂的直方图，仅仅记录了Loss、Accuracy、Graph

import torch
import torch.nn as nn
import torchvision
from torchvision import transforms
from logger import Logger
from tensorboardX import SummaryWriter

# 加载SummaryWriter，设置保存地址。
writer = SummaryWriter('./logs')
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# MNIST dataset
dataset = torchvision.datasets.MNIST(root='./data',
                                     train=True,
                                     transform=transforms.ToTensor(),
                                     download=True)

# Data loader
data_loader = torch.utils.data.DataLoader(dataset=dataset,
                                          batch_size=100,
                                          shuffle=True)


# Fully connected neural network with one hidden layer
class NeuralNet(nn.Module):
    def __init__(self, input_size=784, hidden_size=500, num_classes=10):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out


model = NeuralNet().to(device)

# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.00001)

data_iter = iter(data_loader)
iter_per_epoch = len(data_loader)
total_step = 10000

# Start training
for step in range(total_step):

    # Reset the data_iter
    if (step + 1) % iter_per_epoch == 0:
        data_iter = iter(data_loader)

    # Fetch images and labels
    images, labels = next(data_iter)
    # view作用是将多行tensor拼接为一行，reshape张量形状，如果你不知道你想要多少行，但确定列数，那么你可以将行数设置为-1（同样，不知道多少列，可以将列数设为-1）
    # size获取images的信息（行数，列数）
    images, labels = images.view(images.size(0), -1).to(device), labels.to(device)
    writer.add_graph(model, (images,))
    # Forward pass
    outputs = model(images)
    loss = criterion(outputs, labels)
    writer.add_scalar('Loss', loss, step + 1)

    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Compute accuracy
    _, argmax = torch.max(outputs, 1)
    accuracy = (labels == argmax.squeeze()).float().mean()
    writer.add_scalar('accuracy', accuracy, step + 1)

    if (step + 1) % 100 == 0:
        print('Step [{}/{}], Loss: {:.4f}, Acc: {:.2f}'.format(step + 1, total_step, loss.item(), accuracy.item()))

搞定之后会在文件列表里看到一个logs文件夹，记录就在里面。
在根目录下命令行输入

tensorboard --logidr logs

得到反馈后在浏览器输入"http://localhost:6006"就可以进入tensorboard。

image.png