一、Hook函数概念
Hook函数机制:不改变主体,实现额外功能,像一个挂件,挂钩,hook
1、torch.Tensor.register_hook(hook)
import torch
flag=1
if flag:
w = torch.tensor([1.0], requires_grad=True)
x = torch.tensor([2.], requires_grad=True)
a = torch.add(x, w)
b = torch.add(w, 1)
y = torch.mul(a, b)
a_grad = list()
def grad_hook(grad):
a_grad.append(grad)
a.register_hook(grad_hook)
y.backward()
#查看梯度
# a,b,y为非叶子节点,梯度会被释放
print(w.grad,x.grad,a.grad,b.grad,y.grad)
# # =》tensor([5.]) tensor([2.]) None None None
print(a_grad[0])
# =》tensor([2.])
flage=0
if flag:
w = torch.tensor([1.0], requires_grad=True)
x = torch.tensor([2.], requires_grad=True)
a = torch.add(x, w)
b = torch.add(w, 1)
y = torch.mul(a, b)
a_grad = list()
#改变梯度
def grad_hook(grad):
grad*=2
"""
def grad_hook(grad):
grad*=2
reture 3*grad
"""
w.register_hook(grad_hook)
y.backward()
print(w.grad)
# =》tensor([10.])
同时也可以将改变后的梯度返回。此时会覆盖原来的梯度,如上边"“” “”"内的内容,此时w的梯度会变为原来的6倍(感觉和不返回结果是一样的)
2、torch.nn.Module.register_forward_hook
如要实现下面功能:
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()
self.conv1=nn.Conv2d(1,2,3)
self.pool = nn.MaxPool2d(2,2)
def forward(self,input):
x=self.conv1(input)
y=self.pool(x)
return y
def forward_hook(module,input,output):
fmap_block.append(output)
input_block.append(input)
net=Net()
net.conv1.weight[0].detach().fill_(1)
net.conv1.weight[1].detach().fill_(2)
net.conv1.bias.detach().zero_()
fmap_block,input_block=list(),list()
net.conv1.register_forward_hook(forward_hook)
fake_img=torch.ones((1,1,4,4))
output=net(fake_img)
print("output shape:{}\n output value : {}\n".format(output.shape,output))
print("feature map shape:{}\n output value : {}\n".format(fmap_block[0].shape,fmap_block[0]))
print("input shape:{}\n iinput value:{}".format(input_block[0][0].shape,input_block[0]))
3、torch.nn.Module.register_forward_pre_hook
4、torch.nn.Module.register_backward_hook
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super(Net,self).__init__()
self.conv1=nn.Conv2d(1,2,3)
self.pool = nn.MaxPool2d(2,2)
def forward(self,input):
x=self.conv1(input)
y=self.pool(x)
return y
def forward_pre_hook(module,input):
print("forward_pre_hook_input:{}".format(input))
def backward_hook(module,grad_input,grad_output):
print("backward_hook_input:{}".format(grad_input))
print("backward_hook_output:{}".format(grad_output))
net=Net()
net.conv1.weight[0].detach().fill_(1)
net.conv1.weight[1].detach().fill_(2)
net.conv1.bias.detach().zero_()
fmap_block,input_block=list(),list()
net.conv1.register_forward_pre_hook(forward_pre_hook)
net.conv1.register_backward_hook(backward_hook)
fake_img=torch.ones((1,1,4,4))
output=net(fake_img)
loss_fnc=nn.L1Loss()
target=torch.rand_like(output)
loss=loss_fnc(target,output)
loss.backward()
采用hook函数对卷积层可视化
import torch.nn as nn
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.utils.tensorboard import SummaryWriter
from tools.common_tools import set_seed
import torchvision.models as models
set_seed(1) # 设置随机种子
flag = 1
if flag:
writer = SummaryWriter(comment='test_your_comment', filename_suffix="_test_your_filename_suffix")
# 数据
path_img = "./lena.png" # your path to image
normMean = [0.49139968, 0.48215827, 0.44653124]
normStd = [0.24703233, 0.24348505, 0.26158768]
norm_transform = transforms.Normalize(normMean, normStd)
img_transforms = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
norm_transform
])
img_pil = Image.open(path_img).convert('RGB')
if img_transforms is not None:
img_tensor = img_transforms(img_pil)
img_tensor.unsqueeze_(0) # chw --> bchw
# 模型
alexnet = models.alexnet(pretrained=True)
# 注册hook
fmap_dict = dict()
for name, sub_module in alexnet.named_modules():
if isinstance(sub_module, nn.Conv2d):
key_name = str(sub_module.weight.shape)
fmap_dict.setdefault(key_name, list())
n1, n2 = name.split(".")
def hook_func(m, i, o):
key_name = str(m.weight.shape)
fmap_dict[key_name].append(o)
alexnet._modules[n1]._modules[n2].register_forward_hook(hook_func)
# forward
output = alexnet(img_tensor)
# add image
for layer_name, fmap_list in fmap_dict.items():
fmap = fmap_list[0]
fmap.transpose_(0, 1)
nrow = int(np.sqrt(fmap.shape[0]))
fmap_grid = vutils.make_grid(fmap, normalize=True, scale_each=True, nrow=nrow)
writer.add_image('feature map in {}'.format(layer_name), fmap_grid, global_step=322)
二、CAM and Grad-CAM
CAM局限性:在分类层前需要对通道全局平均池化,这就需要重新训练网络
分析与代码:https://zhuanlan.zhihu.com/p/75894080
import torch.nn as nn
import numpy as np
import torch
from PIL import Image
from torch.autograd import Variable
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
from model.Main_Net import MyEfficientNet
import cv2
def convert_square(img):
w,h=img.size
background = Image.new('RGB', size=(max(w, h), max(w, h)), color=(0, 0, 0)) # 创建背景图,颜色值为0(黑色)
length = int(abs(w - h) // 2) # 一侧需要填充的长度
box = (length, 0) if w < h else (0, length) # 粘贴的位置
background.paste(img, box)
return background
def image_handle(image_path):
transfm=transforms.Compose([
transforms.Grayscale(3),
transforms.Lambda(lambda img:convert_square(img)),
transforms.Resize([456,456]),
transforms.ToTensor(),
])
image = open(image_path, 'rb')
image = Image.open(image)
image = transfm(image)
return image
def load_model(model,model_params,device):
"""
:param cls:分类数量
:param model_params:模型参数
:param device: 是否调用CUDA
:return: 加载好的模型
"""
model_dict = model.state_dict()
pretrained_dict = torch.load(model_params)
pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items() if k[7:] in model_dict}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model = model.to(device)
return model
def comp_class_vec(ouput_vec, index=None):
"""
计算类向量
:param ouput_vec: tensor
:param index: int,指定类别
:return: tensor
"""
if not index:
index = np.argmax(ouput_vec.cpu().data.numpy())
else:
index = np.array(index)
index = index[np.newaxis, np.newaxis]
index = torch.from_numpy(index).cuda()
one_hot = torch.zeros(1, 20,device="cuda").scatter_(1, index, 1)
one_hot.requires_grad = True
class_vec = torch.sum(one_hot * output) # one_hot = 11.8605
return class_vec
def gen_cam(feature_map, grads,img_size):
"""
依据梯度和特征图,生成cam
:param feature_map: np.array, in [C, H, W]
:param grads: np.array, in [C, H, W]
:return: np.array, [H, W]
"""
cam = np.zeros(feature_map.shape[1:], dtype=np.float32) # cam shape (H, W)
weights = np.mean(grads, axis=(1, 2)) #
for i, w in enumerate(weights):
cam += w * feature_map[i, :, :]
cam = np.maximum(cam, 0)# X(cam)和Y(0)逐位进行比较,选择最大值.
cam_min=np.min(cam)
cam_max=np.max(cam)
cam=(cam-cam_min)/(cam_max-cam_min)
cam = cv2.resize(cam, (img_size, img_size))
return cam
def show_cam_on_image(img, mask,fmap_layer_name):
heatmap = cv2.applyColorMap(np.uint8(255*mask), cv2.COLORMAP_JET)
heatmap = np.float32(heatmap) / 255
cam = 0.3*heatmap + np.float32(img)
writer.add_image('feature map in {}'.format(fmap_layer_name), np.uint8(cam*255), dataformats='HWC')
writer=SummaryWriter(log_dir="./logs",flush_secs=60)
model_params=r"./net_params/net.pth"
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
net=MyEfficientNet(20,mode="test")
gragh_inputs=Variable(torch.rand(1,3,456,456))
writer.add_graph(net,(gragh_inputs,))
model=load_model(net,model_params,device)
img = cv2.imread(r"./img.jpg",1)
image=image_handle(r"./img.jpg")
image=image.unsqueeze(0)
image=image.to(device)
fmap_dict=dict()
grad_dict=dict()
for name, sub_module in model.named_modules():
if isinstance(sub_module, nn.Conv2d):
key_name = str(sub_module.weight.shape)
fmap_dict.setdefault(key_name, list())
grad_dict.setdefault(key_name, list())
name_char = name.split(".")
length=len(name_char)
# 定义获取梯度的函数
def backward_hook(module, grad_in, grad_out):
key_name = str(module.weight.shape)
grad_dict[key_name].append(grad_out[0].detach())
# 定义获取特征图的函数
def farward_hook(module, input, output):
key_name = str(module.weight.shape)
fmap_dict[key_name].append(output)
if length == 2:
model._modules[name_char[0]]._modules[name_char[1]].register_forward_hook(farward_hook)
model._modules[name_char[0]]._modules[name_char[1]].register_forward_hook(backward_hook)
if length == 4:
model._modules[name_char[0]]._modules[name_char[1]]._modules[name_char[2]]._modules[name_char[3]].register_forward_hook(farward_hook)
model._modules[name_char[0]]._modules[name_char[1]]._modules[name_char[2]]._modules[name_char[3]].register_forward_hook(backward_hook)
else:
continue
#若只需要单层的,则只需要使用下面的列表存储回调值即可
# grad_block=list()
# fmap_block=list()
# forward
output = model(image)
# backward
model.zero_grad()
class_loss = comp_class_vec(output)
class_loss.backward()
for (fmap_layer_name, fmap_list),(grad_layer_name,grad_list) in list(zip(fmap_dict.items(),grad_dict.items())):
try:
grads_val=grad_list[0].cpu().data.numpy().squeeze()
fmap=fmap_list[0].cpu().data.numpy().squeeze()
cam=gen_cam(fmap,grads_val,456)
img_show=np.float32(cv2.resize(img,(456,456)))/255
show_cam_on_image(img_show,cam,fmap_layer_name)
except:
continue
writer.close()