PyTorch模型训练特征图可视化（TensorboardX）

2019年10月28日 341次阅读来源: Pascal

0、前言

本文所有代码解读均基于PyTroch 1.0，Python3；
本文为原创文章，初次完成于2019.03，最后更新于2019.03；

最近复现的一篇论文一直都难以work，上了特征图可视化后一下子就找到了问题所在，所以今天想梳理一下PyTorch里面的特征图可视化。

大家都知道Tensorflow有一款非常优秀的可视化工具Tensorboard，而PyTorch自身没有可视化功能，但是我们可以寻找替代品，即TensorBoardX。安装过程不多介绍，详见下面的参考链接，里面相应有比较丰富的介绍。

1、Loss可视化

最常见的可视化就是loss曲线作图，这个实现相对比较简单，不多做介绍了

tb_logger.add_scalar('loss_train', loss, curr_step)

《PyTorch模型训练特征图可视化（TensorboardX）》

2、输入图片和标签的可视化

模型不work，第一个应该检查的就是输入输出有没有没给对，因此我们需要将传递给model的 input 和 label 可视化一下。

传递给网络的图片格式往往是 [B,C,H,W] ，范围[0, 1]，数据类型tensor.FloatTensor，但是add_iamge() 能够接受的格式是[C,H,W], 范围[0,1]，数据类型tensor.FloatTensor。

一个是三维的，一个是四维的，这很好解决，我们把每个batch的第一张图拿出来就行了：inptu[0]的形状就是[C,H,W]，符合输入要求。

tb_logger.add_image('image', input[0], curr_step)

tensoroard里面如果出现了猫咪本尊的正确可视化结果，就说明输入图片没问题

《PyTorch模型训练特征图可视化（TensorboardX）》

如果你在Dataloader里对输入图片做了Normalize，显示会出现问题，出现如下所示的乱码，此时需要通过make_grid()函数做一些处理，函数用法具体可见后面的描述。

《PyTorch模型训练特征图可视化（TensorboardX）》

3、单通道特征图的可视化

有时候我们需要把网络内部分节点的特征图可视化出来，这时候上面的方法就不能用了，因为特征图的每个像素点上的数值范围不是[0,1]，而是可正可负，可大可小，因此需要做一些特殊处理。这里就要用到 torchvision.utils.make_grid( )函数，把输入的特征图做一个归一化，把参数normalize设置为True即可，它能帮我们把数据的输入范围调整至[0, 1]之间

def make_grid(tensor, nrow=8, padding=2,
              normalize=False, range=None, scale_each=False, pad_value=0):

这里我把三个中间特征图拼在了一块显示：

from torchvision.utils import make_grid

tb_logger.add_image('feature_map', make_grid([feature_map1, feature_map2, fetare_map3], padding=20, normalize=True, scale_each=True, pad_value=1), curr_step)

需要注意的是：

make_grid() 输入的是Tensor，而不是numpy.ndarray
torchvision.utils.make_grid() 将一组图片绘制到一个窗口，其本质是将一组图片拼接成一张图片

4、多通道特征图的可视化

多通道的特征图的显示和上面的单通道存在一些区别，假设我们从batsh_size=16，channel=20的一个tensor想取出一个多通道特征图可视化，只需要如下操作

feature_map[0].deatch().cpu().unsqueeze(dim=1)

这样就能把一个形状为 [16,20, H, W] 的tesnor取出并转换为 [20, 1, H ,W] 的形状，这与为什么要这么转换，详解第五章节。

完整代码和如下：

tb_logger.add_image('channels', make_grid(feature_map[0].detach().cpu().unsqueeze(dim=1), nrow=5, padding=20, normalize=False, pad_value=1), curr_step)