写了一个有关对卷积网络可视化的小工具,可以直接调用使用,不需要对网络重新训练!【如果用在目标检测网络或者其他网络中需要稍微进行修改】
直接附代码:
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("输入要叠加的原始图像: "))
最终得到的热力图和原始图叠加后:
这样可以更进一步理解卷积的可视化,还有就是可以放在论文中,使自己的论文看着更好看点。放在目标检测中应该需要对上面稍微改一下。