[Pytorch专栏]使用tensorboardX进行训练可视化

叁歲伎倆 2023-06-12 06:19 10阅读 0赞

训练时候可视化loss曲线非常有用，可以很好的观察是否过拟合，还是存在欠拟合，还可以直接观察测试精度

感谢作者开源：[下载链接][Link 1]，直接使用pip安装也是可以的：pip install tensorboardX, tensorflow-gpu

注意安装的时候必须连带tensorflow一起安装，此时会将tensorboard安装好，否则会提示：tensorboard: command not found。tensorboardX借助于tensorboard进行可视化

今天简单讲讲loss曲线的可视化，以maskrcnn-benchmark为例。首先引入一个TensorboardXWriter

from tensorboardX import SummaryWriter
    
    
    class TensorboardXWriter(object):
        """Writes for writing training info to file"""
    
        @staticmethod
        def init(log_dir):
            """
            Args:
                log_dir (str): path to log directory 
            """
            TensorboardXWriter.writer = SummaryWriter(log_dir=log_dir)
    
        @staticmethod
        def write_scalar(names, values, iter):
            #注意传入names和values都是list
            #names表示标题名称，values表示Y轴数值，iter表示迭代次数X轴
            for var_name, var_value in zip(names, values):
                TensorboardXWriter.writer.add_scalar(var_name, var_value, iter)
    
        @staticmethod
        def write_scalars(names, values, iter):
            #注意传入names和values都是list
            #names表示标题名称，values表示Y轴数值，iter表示迭代次数X轴
            for var_name, var_value in zip(names, values):
                TensorboardXWriter.writer.add_scalars(var_name, var_value, iter)
        
        @staticmethod
        def close():
            TensorboardXWriter.writer.close()

这个类包含三个函数，一个是init函数，用于初始化tb\_event目录，还有两个函数write\_scalar和write\_scalars，分别用于加载一个数值和一个列表，但是注意必须封装成list。

value参数如果是PyTorch scalar tensor，则需要调用 .item() 方法获取其数值。一般会使用 add\_scalar 方法来记录训练过程的 loss、accuracy、learning rate 等数值的变化

在tools/train\_net.py中初始化tb\_event

tf_event_dir = os.path.join(output_dir, "tb_events")
    try:
        os.makedirs(tf_event_dir)
    except FileExistsError:
        pass
    
    TensorboardXWriter.init(log_dir=tf_event_dir)

这样完成了初始化操作，然后在maskrcnn\_benchmark/engine/trainer.py中将loss进行可视化

if is_main_process():
        TensorboardXWriter.write_scalars(
            ["train/losses", "train/time"],
            [
                meters.get_metric(
                     metric_names=[
                        "loss",
                        "loss_box_reg",
                        "loss_classifier",
                        "loss_objectness",
                        "loss_rpn_box_reg"
                    ]
                ),
                meters.get_metric(metric_names=["time", "data"]),
            ],
            iteration,
        )

还需要给meters添加get\_metric方法，这个比较简单，关于库函数的导入这里也不罗嗦了，需要特别提醒的一点是必须要进行is\_main\_process的判断，否则在多卡训练的时候会出现KeyError或者FileExistError创建目录失败。这样使用静态函数进行封装的好处是无需进行对象的传递。后台启动tensorboard执行下述语句，--bind\_all可以登录查看服务器上启动的tensorboard，--port 可以自己指定端口号，默认的是6006

> nohup tensorboard --logdir ./ --bind\_all > tb.log 2>&1 &

这样就完成了loss曲线的可视化

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3hpZGlhbnpoaW1lbmc_size_16_color_FFFFFF_t_70][]

使用tensorboardX还可以可视化网络结构或者运行图add\_graph、直方图add\_histogram、 图像add\_image和嵌入向量add\_embedding，可以参考这篇文章了解[详解PyTorch项目使用TensorboardX进行训练可视化][PyTorch_TensorboardX]

固定网络结构使用tensorboardX绘图不如使用netron更加方便，或者使用hiddenlayer更加方便，可惜他们对两阶段的网络支持都不好。如果读者发现好的两阶段网络可视化麻烦留言告知呀

[Link 1]: https://github.com/lanpa/tensorboardX
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3hpZGlhbnpoaW1lbmc_size_16_color_FFFFFF_t_70]: /images/20230531/7e62d11bcdd1498a9d855b7521f6cfa3.png
[PyTorch_TensorboardX]: https://blog.csdn.net/bigbennyguo/article/details/87956434