Pytorch 中使用 Tensorborad 的技巧

---

前言

本文主要记录了如何在 pytorch 中利用 Tensorboard 可视化训练过程

完整详细的 Tensorboard 可视化方法请看第 3 小节，精简版使用可直接看第 4 小节

---

1. 什么是 Tensorboard？

TensorBoard 是一组用于数据可视化的工具。它一开始包含在流行的开源机器学习库 Tensorflow 中，后来拓展到 Pytorch 库中。

TensorBoard 的主要功能包括：

• 可视化模型的网络架构
• 跟踪模型指标，如损失和准确性等
• 检查机器学习工作流程中权重、偏差和其他组件的直方图
• 显示非表格数据，包括图像、文本和音频
• 将高维嵌入投影到低维空间

---

2. 环境准备

2.1. TensorboardX

安装 TensorboardX

pip install tensorboardX

from tensorboardX import SummaryWriter

2.2. Tensorboard

直接使用pytorch当中的自带的Tensorboard，导入所需包的方式如下：

from torch.utils.tensorboard import SummaryWriter # 导入 tensorboard 所需使用的包

本文举例所用的方式为这种

---

3. 记录数据

3.1. 实例化 SummaryWriter

三种实例化方法：

# Creates writer1 object.
# The log will be saved in 'runs/exp'
writer1 = SummaryWriter('runs/exp')

# Creates writer2 object with auto generated file name
# The log directory will be something like 'runs/Aug20-17-20-33'
writer2 = SummaryWriter()

# Creates writer3 object with auto generated file name, the comment will be appended to the filename.
# The log directory will be something like 'runs/Aug20-17-20-33-resnet'
writer3 = SummaryWriter(comment='resnet')

以上展示了三种初始化 SummaryWriter 的方法：

1. 提供一个路径，将使用该 路径 来保存日志（通常为 runs/xxx）
2. 无参数，默认将在当前路径下创建路径 runs/日期时间， 并使用 runs/日期时间 路径来保存日志
3. 提供一个 comment 参数，将使用 runs/日期时间-comment 路径来保存日志

一般情况下，我们对于每次实验新使用第一种方法去实例化一个新的 SummaryWriter

def main(args):

    print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')
    
    # 实例化SummaryWriter对象
    tb_writer = SummaryWriter(log_dir="runs/xxx")

接下来，我们就可以调用 SummaryWriter 实例的各种 add_something 方法向日志中写入不同类型的数据了

想要在浏览器中查看可视化这些数据，只要在命令行中开启 tensorboard 即可：

tensorboard --logdir=<your_log_dir>

• 其中的 既可以是单个 runs 的路径
  • 如上面 writer1 生成的 runs/exp
• 也可以是多个 runs 的父目录
  • 如 runs/ 下面可能会有很多的子文件夹，每个文件夹都代表了一次实验，我们令 --logdir=runs/ 就可以在 Tensorboard 可视化界面中方便地横向比较 runs/ 下不同次实验所得数据的差异

3.2. add_scalar()

使用 add_scalar 方法来记录数字常量

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

各参数含义：

• tag (string)：数据名称，不同名称的数据使用不同曲线展示
• scalar_value (float)：数字常量值
  • 需要注意，这里的 scalar_value 一定是 float 类型，如果是 PyTorch scalar tensor，则需要调用 .item() 方法获取其数值
• global_step (int, optional)：训练的 step
• walltime (float, optional)：记录发生的时间，默认为 time.time()

我们一般会使用 add_scalar 方法来记录训练过程的 loss、accuracy、learning rate 等数值的变化，直观地监控训练过程：

def main(args):
	for epoch in range(args.epochs):
		#####################################################################
		tags = ["train_loss", "accuracy", "learning_rate"]
		tb_writer.add_scalar(tags[0], mean_loss, epoch)
		tb_writer.add_scalar(tags[1], acc, epoch)
		tb_writer.add_scalar(tags[2], optimizer.param_groups[0]["lr"], epoch)
		#####################################################################

训练完每个 epoch 并验证之后进行记录

3.3. add_graph()

使用 add_graph 方法来可视化一个神经网络

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

各参数含义：

• model (torch.nn.Module)：待可视化的网络模型
• input_to_model (torch.Tensor or list of torch.Tensor, optional)：待输入神经网络的变量或一组变量

该方法可以可视化神经网络模型的结构：

def main(args):
	########################################################
	# 将模型写入 tensorboard
    init_img = torch.zeros((1,3,224,224),device=args.device) # (batch, RGB, height, width)
    tb_writer.add_graph(args.model, init_img) 
    ########################################################

需要将 init_img 传入到模型中进行正向传播，会根据输入的数据在模型中正向传播的流程来创建网络结构图

3.4. add_image()

使用 add_image 方法来记录单个图像数据

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

各参数含义：

• tag (string)：数据名称
• img_tensor (torch.Tensor / numpy.array)：图像数据
• global_step (int, optional)：训练的 step
• walltime (float, optional)：记录发生的时间，默认为 time.time()
• dataformats (string, optional)：图像数据的格式
  • 默认为 CHW，即 Channel x Height x Width
  • 还可以是 CHW、HWC 或 HW 等

我们一般会使用 add_image 来实时观察生成式模型的生成效果，或者可视化分割、目标检测的结果，帮助调试模型：

在 train.py 中：

def main(args):
	for epoch in range(args.epochs):
		# add figure into tensorboard
        fig = plot_class_preds(net=model,
                               images_dir="./plot_img",
                               transform=data_transform["val"],
                               num_plot=5,
                               device=device)
        if fig is not None:
            tb_writer.add_figure("predictions",
                                 figure=fig,
                                 global_step=epoch)

在 data_utils.py 中，生成上面所需的 fig

def plot_class_preds(net,
                     images_dir: str,
                     transform,
                     num_plot: int = 5,
                     device="cpu"):  
                     
    """
    :param net: 实例化的模型
    :param images_dir: 待预测图片所在的根目录
    :param transform: 验证集所使用的图像预处理
    :param num_plot: 展示的图片数
    :param device: 使用的 device 信息
	""" 
	
    if not os.path.exists(images_dir):
        print("not found {} path, ignore add figure.".format(images_dir))
        return None

    label_path = os.path.join(images_dir, "label.txt")
    if not os.path.exists(label_path):
        print("not found {} file, ignore add figure".format(label_path))
        return None

    # read class_indict
    json_label_path = './class_indices.json' # 保存类别索引，及其对应的关系
    assert os.path.exists(json_label_path), "not found {}".format(json_label_path)
    json_file = open(json_label_path, 'r')
    # 字典形式载入：{"0": "daisy"}
    flower_class = json.load(json_file)
    # {"daisy": "0"}
    class_indices = dict((v, k) for k, v in flower_class.items())

    # reading label.txt file
    label_info = []
    with open(label_path, "r") as rd:
        for line in rd.readlines():
            line = line.strip()
            if len(line) > 0:
                split_info = [i for i in line.split(" ") if len(i) > 0]
                assert len(split_info) == 2, "label format error, expect file_name and class_name"
                image_name, class_name = split_info
                image_path = os.path.join(images_dir, image_name)
                # 如果文件不存在，则跳过
                if not os.path.exists(image_path):
                    print("not found {}, skip.".format(image_path))
                    continue
                # 如果读取的类别不在给定的类别内，则跳过
                if class_name not in class_indices.keys():
                    print("unrecognized category {}, skip".format(class_name))
                    continue
                label_info.append([image_path, class_name])

    if len(label_info) == 0:
        return None

    # get first num_plot info
    if len(label_info) > num_plot:
        label_info = label_info[:num_plot]

    num_imgs = len(label_info)
    images = []
    labels = []
    for img_path, class_name in label_info:
        # read img
        img = Image.open(img_path).convert("RGB")
        label_index = int(class_indices[class_name])

        # preprocessing
        img = transform(img)
        images.append(img)
        labels.append(label_index)

    # batching images
    images = torch.stack(images, dim=0).to(device)

    # inference
    with torch.no_grad():
        output = net(images)
        probs, preds = torch.max(torch.softmax(output, dim=1), dim=1)
        probs = probs.cpu().numpy()
        preds = preds.cpu().numpy()

    # width, height
    fig = plt.figure(figsize=(num_imgs * 2.5, 3), dpi=100)
    for i in range(num_imgs):
        # 1：子图共1行，num_imgs:子图共num_imgs列，当前绘制第i+1个子图
        ax = fig.add_subplot(1, num_imgs, i+1, xticks=[], yticks=[])

        # CHW -> HWC
        npimg = images[i].cpu().numpy().transpose(1, 2, 0)

        # 将图像还原至标准化之前
        # mean:[0.485, 0.456, 0.406], std:[0.229, 0.224, 0.225]
        npimg = (npimg * [0.229, 0.224, 0.225] + [0.485, 0.456, 0.406]) * 255
        plt.imshow(npimg.astype('uint8'))

        title = "{}, {:.2f}%\n(label: {})".format(
            flower_class[str(preds[i])],  # predict class
            probs[i] * 100,  # predict probability
            flower_class[str(labels[i])]  # true class
        )
        ax.set_title(title, color=("green" if preds[i] == labels[i] else "red"))

    return fig

上述代码所使用的 label.txt 文件内容如下：

1.jpg daisy
2.jpg dandellion
3.jpg roses

上述代码所生成的 .json 文件内容如下：

{
	"0": "daisy"
	"1": "dandelion"
	"2": "roses"
	"3": "sunflowers"
	"4": "tulips"
}

---

4. 快速使用 Tensorboard 可视化信息

本小节主要针对简单使用 add_scalar 函数实现损失函数以及验证信息的可视化

from torch.utils.tensorboard import SummaryWriter  # 导入 tensorboard 所需使用的包

'''在主函数中'''

def main(args, train_loss=None, epoch=None, val_avg_acc=None):
    writer = SummaryWriter(log_dir=args.logdir)
    
    # 在 epoch 循环中
    writer.add_scalar("train_loss", train_loss, epoch)  # 在每次计算完 loss 后添加
    writer.add_scalar("val_acc", val_avg_acc, epoch)

    # 循环结束后
    writer.close()

---

总结

Tensorboard 是非常实用的训练过程可视化的库，本文旨在总结可能会用到的 Tensorboard 的一些可视化方法并给出相关实例。

博文参考