Pytorch可视化模型的特定层

Import特定的库

import torch
import torchvision
from torchvision import transforms
from torchvision import utils
import cv2
import nunmpy as np
from PIL import Image

编写forward_hook函数

该函数用于模型注册，在模型forward时调用, 说白了就是把模型的特定层存到一个字典里

fmap_dict = dict()
module_name = []
def forward_hook(module, input, output):
	key_name = module_name[0]
	fmap_dict[key_name].append(output) # 在模型前向传播时，把key_name层的输出添加到字典中

模型初始化及可视化层注册

model_dir = "保存的模型位置"
net = UNET() #模型结构，模型的那个类
# 注意模型是cpu载入还是gpu载入，在处理输入图片时也应该对应起来
if torch.cuda.is_available():
	net.load_state_dict(torch.load(model_dir))
	net.cuda()
else:
	net.load_state_dict(torch.load(model_dir, map_location='cpu'))

modules_for_plot = (torch.nn.Conv2d) #可以选择可视化层的类型可视化
for name, sub_module in net.named_modules():
	if isinstance(sub_module, modules_for_plot):
		if name == 'layer_name': # 填写可视化层的名字，上一篇博客有写到怎么看每一层名字
			module_name.append(name)
			fmap_dict.setdefault(name,list())
			sub_module.register_forward_hook(forward_hook)

输入图片进行处理并进行可视化

path_img = "可视化图片的路径"

img_transform = transforms.Compose([
   transforms.Resize((640, 640)),
 transforms.ToTensor()
]) #与网络训练保持一致
img = Image.open(path_img).convert('RGB')
if img_transform is not None:
    img_tensor = img_transform(img)*255.0
    img_tensor.unsqueeze_(0)
    img_tensor = img_tensor.cuda()

output = net(img_tensor)

for layer_name, fmap_list in fmap_dict.items():
    if fmap_list is None:
        continue
    fmap = fmap_list[0]
    if fmap.shape[0] < fmap.shape[1]:
        fmap.transpose_(0,1)
    fmap = fmap.sigmoid()
    fmap = F.interpolate(fmap, size=[512,1024], mode='bilinear')#可视化层的特征层数过多导致显存不够时，调小尺寸
    fmap_grid = utils.make_grid(fmap, normalize=True, scale_each=True, nrow=6, pad_value=255) #将特征层拆分成格子显示
    fmap_grid = fmap_grid.permute(1, 2, 0)

    cv2.imwrite("vis.jpg", np.array(fmap_grid.cpu().detach().numpy()*255, dtype=np.uint8))