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

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

感谢作者开源：下载链接，直接使用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曲线的可视化

使用tensorboardX还可以可视化网络结构或者运行图add_graph、直方图add_histogram、 图像add_image和嵌入向量add_embedding，可以参考这篇文章了解详解PyTorch项目使用TensorboardX进行训练可视化

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