import torch
from torchvision import models, transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
#import scipy.misc
# 获取第k层的特征图
def get_k_layer_feature_map(feature_extractor, k, x):
#打印网络的层数
print(feature_extractor)
with torch.no_grad():
for index,layer in enumerate(feature_extractor):
x = layer(x)
if k == index:
print(layer)
return x,str(layer)
# 可视化特征图
def show_feature_map(feature_map,layer,k):
num = 0
for i in range(len(layer)):
if layer[i]!='(':
num+=1
else:
break
layer_ = layer[:num]
num = 0
feature_map = feature_map.squeeze(0)
feature_map = feature_map.cpu().numpy()
print(feature_map.shape)
feature_map_num = feature_map.shape[0]
row_num = np.ceil(np.sqrt(feature_map_num))
plt.figure()
for index in range(1, feature_map_num+1):
plt.subplot(1, 1, 1)
plt.imshow(feature_map[index-1], cmap='gray')
plt.savefig('./feature_map/{}-{}_channel{}_feature_map.jpg'.format(k,layer_,index))
#plt.axis('off')
#scipy.misc.imsave(str(index)+".png", feature_map[index-1])
# plt.savefig('./feature_map.jpg')
# plt.show()
if __name__ == '__main__':
# 初始化图像的路径
image_dir = r"husky.png"
# 定义提取第几层的feature map
k = 1
# 导入模型
model = vgg(model_name='vgg16',num_classes=1000).to(device) #模型。需修改
#读取图像信息
image_info = #这里需要放入的信息为原始图像经过dataloader预处理后的结果(也就是真正放入模型的数据),需要根据不同网络手动修改。
#权重路径
weights_path = './vgg16-397923af.pth' #推理权重路径,需修改
model.load_state_dict(torch.load(weights_path,map_location=device))
model.eval()
model = model.cuda()
image_info = image_info.cuda()
# alexnet只有features部分有特征图
# classifier部分的feature map是向量
feature_extractor = model.features
feature_map ,layer= get_k_layer_feature_map(feature_extractor, k, image_info)
show_feature_map(feature_map,layer,k)
该脚本可以输出网络层内容,以及想要的某一层的全部通道上的特征图。细节可根据代码进行修改。
原始内容参考:(110条消息) Pytorch从入门到放弃(8)——可视化FeatureMap(特征图)_澎湃de家夥的博客-CSDN博客_feature map可视化 pytorch