PyTorch可视化-在PyTorch中使用TensorBoard

在训练模型的过程中，可以使用TensorBoard对工作流中的一些数据进行可视化。 TensorBoard能够跟踪实验指标，例如损失和准确性，可视化模型图，将嵌入物投影到较低维度的空间等等。

PyTorch在torch.utils.tensorboard中的对tensorboard进行了封装，我们可以通过pytorch提供的接口使用tensorboard。

安装

需要安装pytorch和tensorboard。

pip install tensorboard

可视化

当使用tensorboard生成可视化文件之后，即可使用以下命令进行读取。

可在浏览器对其查看，运行下面命令，在浏览器访问输出的地址http://localhost:6007/：

tensorboard --logdir=runs

也可在notebook可以使用插件，可视化结果会输出在cell下方，有时可能会加载比较慢：

%load_ext tensorboard
%tensorboard --logdir=runs

其中runs是默认存储的路径，后面会讲。

SummaryWriter

torch.utils.tensorboard.writer.SummaryWriter(log_dir=None, comment=’’, purge_step=None, max_queue=10, flush_secs=120, filename_suffix=’’)

log_dir

SummaryWriter通过log_dir参数指定保存路径：

# https://github.com/pytorch/pytorch/blob/master/torch/utils/tensorboard/writer.py#L205
        if not log_dir:
            import socket
            from datetime import datetime
            current_time = datetime.now().strftime('%b%d_%H-%M-%S')
            log_dir = os.path.join(
                'runs', current_time + '_' + socket.gethostname() + comment)
        self.log_dir = log_dir

add_scalar

add_scalar(tag, scalar_value, global_step=None, walltime=None)

import torch
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter()

for i in range(10):
    writer.add_scalar('y=2x', 2*i, i)
writer.close()

这段代码我运行了3次，在当前路径下生成runs文件夹，并对应生成3个子文件夹，里面存的就是记录的数据events.out.tfevents.xxx：

在tensorboard中可以看到渲染出3条重叠的线：

log_dir自定义保存的路径：

import torch
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir='my_runs')

for i in range(10):
    writer.add_scalar('y=2x', 2*i, i)
writer.close()

同样运行3次，在当前路径下生成my_runs文件夹，并保存了3个文件events.out.tfevents.xxx：

此时在tensorboard中只能看到一条线：

可见每个文件夹记录一条线，或者扩展来说，每个文件夹存储一组数据。
如果需要保存成不同的线。则需要更改文件夹名称，比如：log_dir=my_runs1、log_dir=my_runs2、log_dir=my_runs3；或者：log_dir=my_runs/exp1、log_dir=my_runs/exp2、log_dir=my_runs/exp3

add_scalars

使用dict向summary添加多个标量数据，利用tag将多个不同数据添加到同一个标量图中。

add_scalars(main_tag, tag_scalar_dict, global_step=None, walltime=None)

from torch.utils.tensorboard import SummaryWriter
import numpy as np

writer = SummaryWriter()
r = 5
for i in range(10):
    writer.add_scalars('run_14h', {
     'xsinx':i*np.sin(i/r),
                                    'xcosx':i*np.cos(i/r),
                                    'tanx': np.tan(i/r)}, i)
writer.close()

add_image

添加图像数据到summary，可在IMAGES栏查看

add_image(tag, img_tensor, global_step=None, walltime=None, dataformats=‘CHW’)

from torch.utils.tensorboard import SummaryWriter
import numpy as np
img = np.zeros((3, 100, 100))
img[0] = np.arange(0, 10000).reshape(100, 100) / 10000
img[1] = 1 - np.arange(0, 10000).reshape(100, 100) / 10000

img_HWC = np.zeros((100, 100, 3))
img_HWC[:, :, 0] = np.arange(0, 10000).reshape(100, 100) / 10000
img_HWC[:, :, 1] = 1 - np.arange(0, 10000).reshape(100, 100) / 10000

writer = SummaryWriter()
writer.add_image('my_image', img, 0)

# 默认的dataformats是'CHW'，如果是非默认维度，可以设置dataformats参数。
writer.add_image('my_image_HWC', img_HWC, 0, dataformats='HWC')
writer.close()

如果需要多张图片，可以使用torchvision.utils.make_grid()将多张图片转为一张；也可以使用下面的add_images。

add_images

添加图像数据到summary

add_images(tag, img_tensor, global_step=None, walltime=None, dataformats=‘NCHW’)

from torch.utils.tensorboard import SummaryWriter
import numpy as np

img_batch = np.zeros((16, 3, 100, 100))
for i in range(16):
    img_batch[i, 0] = np.arange(0, 10000).reshape(100, 100) / 10000 / 16 * i
    img_batch[i, 1] = (1 - np.arange(0, 10000).reshape(100, 100) / 10000) / 16 * i

writer = SummaryWriter()
writer.add_images('my_image_batch', img_batch, 0)
writer.close()

也可以将matplotlib的figure作为图像添加到summary。

add_figure

add_figure(tag, figure, global_step=None, close=True, walltime=None)

import torch
from torch.utils.tensorboard import SummaryWriter
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

writer = SummaryWriter()

figure = plt.figure()
x = np.linspace(0, 3*np.pi, 500)
plt.plot(x, np.sin(x**2))
plt.title('chirp')

writer.add_figure('figure/chirp', figure)

plt.close()
writer.close()

add_histogram

add_histogram(tag, values, global_step=None, bins=‘tensorflow’, walltime=None, max_bins=None)

import torch
from torch.utils.tensorboard import SummaryWriter
import numpy as np

writer = SummaryWriter()

x = np.array([j for j in range(10)])
writer.add_histogram('histogram', x, i)
writer.close()

add_pr_curve

添加PR曲线。绘制PR曲线可以让你了解模型在不同阈值设置下的性能。

add_pr_curve(tag, labels, predictions, global_step=None, num_thresholds=127, weights=None, walltime=None)

为每个目标提供ground truth标签和预测置信度（通常是模型的输出）

from torch.utils.tensorboard import SummaryWriter
import numpy as np


writer = SummaryWriter()

labels = np.random.randint(2, size=100)  # 随机生成的伪标签
predictions = np.random.rand(100)  # 随机生成的伪预测置信度

writer.add_pr_curve('pr_curve', labels, predictions, global_step=0)
writer.close()

add_hparams

add_hparams(hparam_dict, metric_dict, hparam_domain_discrete=None, run_name=None)

from torch.utils.tensorboard import SummaryWriter
with SummaryWriter() as w:
    for i in range(5):
        w.add_hparams({
     'lr': 0.1*i, 'bsize': i},
                      {
     'hparam/accuracy': 10*i, 'hparam/loss': 10*i})

add_graph

add_graph(model, input_to_model=None, verbose=False)

import torch
import torchvision
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

# Writer will output to ./runs/ directory by default
writer = SummaryWriter()

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
trainset = datasets.MNIST('data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
model = torchvision.models.resnet50(False)
# Have ResNet model take in grayscale rather than RGB
model.conv1 = torch.nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
images, labels = next(iter(trainloader))

grid = torchvision.utils.make_grid(images)
writer.add_image('images', grid, 0)
writer.add_graph(model, images)
writer.close()

…
还有其他一些可以可视化的，觉得很少用到，不列举了

总结

TensorBoard通过调用不同的函数添加对应的UI组件，比如add_scalar会添加scalar栏；
通过不同的tag添加不同的显示区域，相同tag一起展示；
通过不同的文件夹区分不同的数据组，比如相同tag的标量，通过不同文件夹区分；这样就可以将同一算法不同训练参数得到的指标放在一起展示做比较。

其他

PyTorch Docs
PyTorch github
以 DataFrame 格式访问 TensorBoard 数据