Tensorboard & SummaryWriter & PyTorch

Tensorboard & SummaryWriter & PyTorch

TensorBoard 是 2015年Google 推出的神经网络可视化开源服务包，可以记录每次训练的摘要(Summary)。
在Pytorch中的API是torch.utils.tensorboard，其中SummaryWriter类是记录数据的主要入口。

SummaryWriter

SummaryWriter(log_dir=None, comment='', purge_step=None, max_queue=10, flush_secs=120, filename_suffix='')
该类提供一个接口，在给定文件夹下创建事件文件(event file)，并将训练过程中的数据实时地添加到该事件文件中。
参数列表：

• log_dir：设置保存的文件路径位置，默认创建并保存在runs文件夹下，默认命名格式为CURRENT_DATETIME_HOSTNAME
• comment：给log_dir设置个后缀名，默认为空。tips：可以只设置comment参数（不设置log_dir参数），给每次训练的event file文件添加标签，便于自己辨别。
• purge_step：详情点此
• max_queue：详情点此
• flush_secs：详情点此
• filename_suffix：详情点此

SummaryWriter类的方法

1. add_scalar：将标签数据添加到Summary中。
   <参数列表>$\downarrow$
   (tag, scalar_value, global_step=None, walltime=None, new_style=False, double_precision=False)，详情点此
2. add_scalars：将多个标签数据添加到Summary中。
   <参数列表>$\downarrow$
   (main_tag, tag_scalar_dict, global_step=None, walltime=None)，详情点此
3. add_graph：将网络图添加到Summary中。
   <参数列表>$\downarrow$
   (model, input_to_model=None, verbose=False, use_strict_trace=True)，详情点此
4. add_image：将图像数据添加到Summary中。
   <参数列表>$\downarrow$
   (tag, img_tensor, global_step=None, walltime=None, dataformats=‘CHW’)，详情点此
5. add_images：将批量图像数据添加到Summary中。
   <参数列表>$\downarrow$
   (tag, img_tensor, global_step=None, walltime=None, dataformats=‘NCHW’)，详情点此

实例

import torch
import torch.nn as nn
import torch.utils.data
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms


# 随便搭建个网络
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=4, stride=4)  # 7 * 7
        self.bn = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)  # 7 * 7
        self.relu = nn.ReLU()
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)  # 4 * 4
        self.fc = nn.Linear(128, 10, bias=True)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.conv3(x)
        x = torch.mean(x, [-2, -1])
        result = self.fc(x)
        return result


transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_set = datasets.MNIST('mnist_train', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32, shuffle=True)
model = MyModel()
###################################################
# 给默认生成的文件添加一个后缀名
writer = SummaryWriter(comment='.test1')  # 实际文件名为May14_15-49-55_DESKTOP-U3KRP5F.test1
images, labels = next(iter(train_loader))
# 将网络图添加到Summary中，送进去一批图片，可以查看数据的处理过程
writer.add_graph(model, images)
# writer.add_image("image", images)
image = images[0]
writer.add_image("image", image, dataformats='CHW')  # 默认格式是(N, 3, H, W),而手写字是一维图像，需要设置dataformats参数
writer.add_images("images", images, dataformats='NCHW')
###################################################

loss_function = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
epochs = 5
for epoch in range(epochs):
    model.train()
    acc_num, accu_loss = 0, 0
    for step, data in enumerate(train_loader):
        images, labels = data
        predicts = model(images)
        predict_labels = torch.max(predicts, dim=1)[1]
        acc_num += torch.eq(predict_labels, labels).sum()

        optimizer.zero_grad()
        loss = loss_function(predicts, labels)
        loss.backward()
        optimizer.step()
        accu_loss += loss.item()

        if (step+1) % 100 == 0 or (step+1) == len(train_loader):
            print("epoch: {}, [{}/{}], acc: {:.4f}, loss: {:.4f}".format(epoch, (step+1), len(train_loader),
                                                                         acc_num / ((step + 1) * 32),
                                                                         accu_loss / (step + 1)))
    train_acc = acc_num / (len(train_loader) * 32)
    loss = accu_loss / len(train_loader)

#######################################################################
    # 将‘train_acc’的标签数据（曲线图）添加到Summary中。
    writer.add_scalar('train_acc', train_acc, epoch)
    # 同上，添加多个标签数据
    writer.add_scalars('train_acc and loss', {'train_acc': train_acc, 'loss': loss}, epoch)
#######################################################################

在终端切换到当前工作目录下，输入下面指令打开tensorboard：

tensorboard --logdir=./runs

注意：./runs是默认创建的文件夹，如果在实例化SummaryWriter是创建了自己的文件夹，记得上面指令改一下路径。

结果

add_scalar/add_scalars:

add_image/add_images:

add_graph: