分类卷积网络的可视化【附代码】

写了一个有关对卷积网络可视化的小工具，可以直接调用使用，不需要对网络重新训练！【如果用在目标检测网络或者其他网络中需要稍微进行修改】

直接附代码：

import cv2
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn


def Cnn_View(cnn_output, classes_layer, Or_img):
    '''
    cnn_output.size(1)是获得上一层的通道数
    如果你用的是CPU推理的，那么在cam处你应该将张量放在cpu上【我这里默认用的cuda】
    因为我的网络输入大小为224*224大小，所以需要对resize成224*224，以保证叠加图像大小一致！！
    最后将热力图和原始图进行一个叠加
    本代码是放在平均池化之后，分类层之前的，即classes_layer是定义的全连接层,如果放在其他层可以把该层注释掉
    '''
    Or_img = cv2.imread(Or_img)
    cam = nn.Conv2d(cnn_output.size(1), 1, 1, 1, 1).cuda()  # 512是上一层卷积输出后的通道，可以根据自己的网络修改
    cnn_output1 = cnn_output

    cnn_output = torch.flatten(cnn_output, 1)  # 平铺  (batch_size,512*7*7)
    cnn_output = classes_layer(cnn_output)  # 分类  （batch_size,512*7*7，num_classes）
    preds = torch.softmax(cnn_output[0], dim=-1).cpu().numpy()

    cam_output = cam(cnn_output1)

    cam_output[0, :, :, 1] = cam_output[0, :, :, 1 if preds.all() > 0.5 else 0]
    
    cam_output /= 10
    cam_output[cam_output < 0] = 0
    cam_output[cam_output > 1] = 1
    cam_output = cam_output.cpu().numpy()
    cam_output = cam_output.squeeze(0)
    img = cam_output.transpose(1, 2, 0)
    img = cv2.resize(img, (224, 224))
    img = np.uint8(255 * img)
    heatmap = cv2.applyColorMap(img, cv2.COLORMAP_JET)
    Or_img = cv2.resize(Or_img, (224, 224))
    out = cv2.addWeighted(Or_img, 0.8, heatmap, 0.4, 0)
    plt.axis('off')
    plt.imshow(out[:, :, ::-1])
    plt.show()

在网络中调用该函数，放在网络的forward(self,x)函数下，如果是分类网络，可以放在平均池化层后面 

例如：

    def forward(self, x):

        x = self.features(x)  # 主干网络  14*14输出通道512
        x = self.avgpool(x)  # 平均池化  7*7输出通道512   (batch_size,512,7,7)
        Cnn_View(x, self.classifier, input("输入要叠加的原始图像: "))

---

最终得到的热力图和原始图叠加后：

---

这样可以更进一步理解卷积的可视化，还有就是可以放在论文中，使自己的论文看着更好看点。放在目标检测中应该需要对上面稍微改一下。