[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曲线的可视化

[Pytorch专栏]使用tensorboardX进行训练可视化_第1张图片

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

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

你可能感兴趣的:(Pytorch专栏)