PyTorch学习笔记-TensorBoard

1. TensorBoard的安装与绘图

TensorBoard 原本是 TensorFlow 的可视化工具，PyTorch 从1.2.0开始支持 TensorBoard。之前的版本也可以使用 TensorBoardX 代替。

先进入 Anaconda 的 PyTorch 环境，安装 TensorBoard：

pip install tensorboard

在项目中新建一个文件夹 logs：

TensorBoard 的工作流程简单来说是将代码运行过程中的，某些你关心的数据保存在这个文件夹中（由代码中的 writer 完成），再读取这个文件夹中的数据，用浏览器显示出来（在命令行运行 TensorBoard 完成）。

我们先绘制一个 y = x 的图像，运行以下代码：

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('logs')

for x in range(100):
    writer.add_scalar('y=x', x, x)  # tag='y=x', scalar_value=x, global_step=x

writer.close()

首先来看函数 add_scalar，该函数主要有三个参数：

tag 表示图像的标题；scalar_value 表示 y 的值；global_step 表示 x 的值。

运行后会看到 logs 文件夹下生成了一个文件：

然后我们在 PyCharm 终端的 PyTorch 环境中打开 TensorBoard：

tensorboard --logdir logs

打开 http://localhost:6006/ 即可看到绘制的图像：

如果因为某些原因导致端口冲突可以指定端口：

tensorboard --logdir logs --port 6007

2. TensorBoard添加图像

现在来看一下函数 add_image，该函数也是主要有三个参数：

可以看到传入的图片数据类型需要是 torch.Tensor 或 numpy.ndarray，因此我们需要先安装一下 OpenCV 库，还是在 PyTorch 环境中安装：

pip install opencv-python

使用 PIL 打开一个图像，将其转换成 Numpy 数组：

from PIL import Image
import numpy as np

img_path = 'dataset/hymenoptera_data/train/ants_image/0013035.jpg'
img_PIL = Image.open(img_path)
img_array = np.array(img_PIL)
print(type(img_array))  # 
print(img_array.shape)  # (512, 768, 3)

可以看到图片的形状是三维的数据，前两个数据分别表示高度和宽度，第三个数据表示通道数，add_image 函数传入图片时有一些规定：

意思是默认的图像格式为：(3, H, W)，如果要改为 (H, W, 3) 的话需要设置 dataformats 参数：

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

writer = SummaryWriter('logs')
img_path = 'dataset/hymenoptera_data/train/ants_image/0013035.jpg'
img_PIL = Image.open(img_path)
img_array = np.array(img_PIL)

writer.add_image('img_test', img_array, 1, dataformats='HWC')

writer.close()

运行后打开 TensorBoard，在 IMAGES 页面下可以看到图片：