Pytorch可视化语义分割特征图

简介

深度学习中有时候可视化特征图是必要的，特别是对于语义分割任务，合理分析特征图也许能够发现新的idea!接下来讲解一种Pytorch框架下的可视化方法，这里采取的网络模型为Deeplabv3+，首先介绍一些背景知识和几个函数的使用方法。通常网络模型中的特征图的shape为$[n,c,h,w]$，分别代表batchsize, channel, height, width. 换句话说，我们需要处理的特征图实际上是四维度的Tensor，考虑到可视化特征图需要保存目标图像(二维)，因此调用训练好的模型测试时应当设置batchsize为1，则可假定待可视化特征图的shape为$[1,c,h,w]$。这里的可视化思路取决于个人，但是最常见的几个想法应该是：(1)单独拿出一个channel进行可视化；(2)在通道维度上取最大值得到一个channel进行可视化；(3)在通道维度上取平均值得到一个channel进行可视化。那么这里就需要用一些函数来对指定维度进行操作，可以对GPU上的Tensor直接操作，涉及到torch.max()，torch.mean()等函数；也可以将Tensor转到cpu上用numpy库处理，涉及到np.max()，np.mean()等函数。这里简单介绍下这些函数，建议直接查Numpy官方文档和Pytorch官方文档。

函数使用

1. torch.max(input, dim, keepdim=False, *, out=None) -> (Tensor, LongTensor)
   Returns a namedtuple (values, indices) where values is the maximum value of each row of the input tensor in the given dimension dim. And indices is the index location of each maximum value found (argmax).If keepdim is True, the output tensors are of the same size as input except in the dimension dim where they are of size 1. Otherwise, dim is squeezed (see torch.squeeze()), resulting in the output tensors having 1 fewer dimension than input.
   文档指出：返回的是tuple，包含values和indices两项，且维度默认是会被压缩一维(可以设置不压缩)。假定待处理的特征图Tensor的shape为$[1,c,h,w]$，且该Tensor变量名为x：x_channel_max,index = torch.max(x,dim=1)即为使用方法，x_channel_max为通道维度取最大值返回的values，其shape为$[1,h,w](被压缩一维)$，1为batchsize；可见dim=1指定了对1维度进行求max操作，$1,c,h,w$分别对应的维度为$0,1,2,3$；dim=1就是在通道维度上取最大值，因此对于上图中的toch.max(a,1)也是对1维度操作，原张量a的shape为[h,w]，对1维度操作就是对w维度操作，即求行最大值。
2. torch.mean(input, dim, keepdim=False, *, out=None) → Tensor
   理解同上，返回值只有一个，不是tuple，而是指定维度上的平均值。
3. numpy.amax(a, axis=None, out=None, keepdims=, initial=, where=)[source]
   理解同上，axis用来指定维度，与numpy.max()是等同函数，其他不再赘述。

可视化特征图

Deeplabv3+的网络结构如上图，在输入图像是$[3,513,513]$的设置下，假定这里想可视化上图中红色虚线框住的浅绿色特征图$X$，该特征图是ASPP模块后经过$1\times 1$卷积的输出，其shape为$[1,256,33,33]$。可以在Decoder的代码中找到该特征图X，考虑到该特征图比较小，因此可以考虑上采样至$[1,256,513,513]$的大小(费时间)，再获取空间维度最大值(或均值)；也可以在获取空间维度最大值(或均值)之后得到$[1,33,33]$，再上采样至$[1,513,513]$(时间短)。下面提供一种可视化做法(只有代码片段，仅供参考)。
(1)这是Decoder的代码片段：

    def forward(self, x, low_level_feat):
        low_level_feat = self.conv1(low_level_feat)
        low_level_feat = self.bn1(low_level_feat)
        low_level_feat = self.relu(low_level_feat) #这里1*1卷积得到上图天蓝色虚线框住的特征图 L   
        x_visualize = x #获取上图红色虚线框住的特征图 X，shape为[1,256,33,33]
        x_visualize = F.interpolate(x_visualize, size=(513,513), mode='bilinear', align_corners=False)#这种做法费时间 shape为[1,256,513,513]
        '''
        #这个做法省时间
        x_visualize,index = torch.max(x,dim = 1) #shape为[1,33,33]
        x_visualize = F.interpolate(x_visualize, size=(513,513), mode='bilinear', align_corners=False) #shape为[1,513,513]
		'''
        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=False)
        x = torch.cat((x, low_level_feat), dim=1)
        x = self.last_conv(x)
        return x,x_visualize #返回想要可视化的特征图给输出端

(2)这是主函数测试时代码片段：

       output,x_visualize = self.model(image) #测试时，获取网络模型返回的想要可视化的特征图
       x_visualize = x_visualize.cpu().numpy() #用Numpy处理返回的[1,256,513,513]特征图
       x_visualize = np.max(x_visualize,axis=1).reshape(513,513) #shape为[513,513]，二维
       x_visualize = (((x_visualize - np.min(x_visualize))/(np.max(x_visualize)-np.min(x_visualize)))*255).astype(np.uint8) #归一化并映射到0-255的整数，方便伪彩色化
       savedir =  '/home/mfx/xmf/Seg/pytorch-deeplab-xception-master/run/pascal/deeplab-resnet/'        
       if not os.path.exists(savedir+'val_pred_temp'):
              os.mkdir(savedir+'val_pred_temp')   
       x_visualize = cv2.applyColorMap(x_visualize, cv2.COLORMAP_JET)  # 伪彩色处理   
       cv2.imwrite(savedir+'val_pred_temp/'+str(i)+'.jpg',x_visualize) #保存可视化图像

可视化效果如下，第一幅是原图(需要被裁剪)，第二幅是groundtruth，第三幅是特征图X

如果采取的是np.mean()函数，则效果为如下。对比np.max()函数的效果可以发现，np.max()得到的特征图边界更加明显，而np.mean()函数得到的特征图边界较为模糊，但语义上的类别一致性似乎更好。

同理：如果想可视化网络结构图中天蓝色虚线框住的特征图 L ，其shape为[1,48,129,129]，只需稍稍修改Decoder的代码片段为：

    def forward(self, x, low_level_feat):
        low_level_feat = self.conv1(low_level_feat)
        low_level_feat = self.bn1(low_level_feat)
        low_level_feat = self.relu(low_level_feat) #这里1*1卷积得到上图天蓝色虚线框住的特征图 L   
        x_visualize = low_level_feat #获取上图天蓝色虚线框住的特征图 L，shape为[1,48,129,129]
        x_visualize = F.interpolate(x_visualize, size=(513,513), mode='bilinear', align_corners=False)#这种做法费时间 shape为[1,48,513,513]
        '''
        #这个做法省时间
        x_visualize,index = torch.max(low_level_feat,dim = 1) #shape为[1,129,129]
        x_visualize = F.interpolate(x_visualize, size=(513,513), mode='bilinear', align_corners=False) #shape为[1,513,513]
		'''
        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=False)
        x = torch.cat((x, low_level_feat), dim=1)
        x = self.last_conv(x)
        return x,x_visualize #返回想要可视化的特征图给输出端

可以看出，low_level_feat对应的低层次高分辨率特征图果然是细节信息较大，但是却不具备语义信息，所以语义分割网络才会融合低层次和高层次特征，但是可以发现低层次特征存在着大量噪声(无用信息)，如同文献描述一样。

总结

可视化特征图的策略有很多，根据自己的思路进行合理的分析说不定会带来新的想法。文献中经常提到，语义分割中的高层次特征图分辨率低，但语义性强；而低层次特征图分辨率高，细节丰富。通过以上的Pytorch可视化特征图，更深一步理解了语义分割任务。

补充

考虑到大家想问完整的代码，这里重新组织了下。feats即为自己想可视化的特征，传入feature_vis()函数即可。

import torch
import torch.nn.functional as F
import cv2
def feature_vis(feats): # feaats形状: [b,c,h,w]
     output_shape = (512,1024) # 输出形状
     channel_mean = torch.mean(feats,dim=1,keepdim=True) # channel_max,_ = torch.max(feats,dim=1,keepdim=True)
     channel_mean = F.interpolate(channel_mean, size=output_shape, mode='bilinear', align_corners=False)
     channel_mean = channel_mean.squeeze(0).squeeze(0).cpu().numpy() # 四维压缩为二维
     channel_mean = (((channel_mean - np.min(channel_mean))/(np.max(channel_mean)-np.min(channel_mean)))*255).astype(np.uint8)
     savedir = '/home/mfx/xmf/mmsegmentation-master/work_dirs/'
     if not os.path.exists(savedir+'feature_vis'): os.makedirs(savedir+'feature_vis') 
     channel_mean = cv2.applyColorMap(channel_mean, cv2.COLORMAP_JET)
     cv2.imwrite(savedir+'feature_vis/'+ '0.png',channel_mean)