Pytorch的hook函数

hook函数是勾子函数，用于在不改变原始模型结构的情况下，注入一些新的代码用于调试和检验模型，常见的用法有保留非叶子结点的梯度数据（Pytorch的非叶子节点的梯度数据在计算完毕之后就会被删除，访问的时候会显示为None），又或者查看模型的层与层之间的数据传递情况（数据维度、数据大小等），抑或是在不修改原始模型代码的基础上可视化各个卷积特征图。

Pytorch提供了四种hook函数

  1. torch.tensor.register_hook(hooc_func)
  2. torch.nn.Module.register_forward_hook(hook_func)
  3. torch.nn.Module.register_forward_pre_hook(hook_func)
  4. torch.nn.Module.register_backward_hook

1. torch.tensor.register_hook(hooc_func)

解释：注册一个反向传播hook函数，其函数签名如下

def hook(grad):...

输入参数为张量的梯度，实现的hook函数可以在此修改梯度数据（原地修改或者通过返回值返回），或者在此将梯度数据保存、裁剪等。

示例 1

# leaf node data
x = torch.Tensor([0, 1, 2, 3]).requires_grad_()
y = torch.Tensor([4, 5, 6, 7]).requires_grad_()
w = torch.Tensor([1, 2, 3, 4]).requires_grad_()# intermediate variable
z = x + y# output 
o = torch.dot(w, z)# backward to calculate gradient
o.backward()# print gradient infomation
print('x.grad:', x.grad) # tensor([1., 2., 3., 4.])
print('y.grad:', y.grad) # tensor([1., 2., 3., 4.])
print('w.grad:', w.grad) # tensor([ 4.,  6.,  8., 10.])
print('z.grad:', z.grad) # None
print('o.grad:', o.grad) # None

输出：

x.grad: tensor([1., 2., 3., 4.])
y.grad: tensor([1., 2., 3., 4.])
w.grad: tensor([ 4.,  6.,  8., 10.])
z.grad: None
o.grad: None

可以看到代码中的非叶子节点z, o的梯度信息(grad)在计算之后立即被释放，因此都等于None，如果需要显式地声明需要保留非叶子节点的grad，需要使用retain_grad方法，如下例：

import torch 
a = torch.ones(5)
a.requires_grad = Trueb = 2*ab.retain_grad()   # 让非叶子节点b的梯度保持
c = b.mean()
c.backward()print(f'a.grad = {a.grad}\nb.grad = {b.grad}')

输出：

a.grad = tensor([0.4000, 0.4000, 0.4000, 0.4000, 0.4000])
b.grad = tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])

retain_grad()方法会增加显存的占用，我们可以使用hook获取梯度信息而不需要显式地使用retain_grad()强制系统保存梯度信息，如下例：

import torcha = torch.ones(5).requires_grad_()b = 2 * aa.register_hook(lambda x:print(f'a.grad = {x}'))
b.register_hook(lambda x: print(f'b.grad = {x}'))  c = b.mean()print('begin backward'.center(30, '-'))
c.backward()
print('end backward'.center(30, '-'))

输出：

--------begin backward--------
b.grad = tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])
a.grad = tensor([0.4000, 0.4000, 0.4000, 0.4000, 0.4000])
---------end backward---------

上述例子中我们使用hook对tensor的grad进行访问，没有使用retain_grad对信息进行保存。输出结果表明，hook执行的时间是在backward之间，从后往前依次执行，首先输出b的grad，然后输出a的grad，最后结束backward过程。

上述过程都没有对梯度信息进行改变，其实，如果hook函数的有返回值或者将输入参数grad原地进行修改的话，那么之后的梯度信息都会被改变，这一机制简直就是为梯度裁剪量身定制的。

如下例：

import torchdef hook(grad):torch.clamp_(grad, min=0.5, max=0.2)print(grad)a = torch.ones(5).requires_grad_()
b = 2 * aa.register_hook(hook)
b.register_hook(hook)  c = b.mean()print('begin backward'.center(30, '-'))
c.backward()
print('end backward'.center(30, '-'))

输出：

--------begin backward--------
tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])
tensor([0.2000, 0.2000, 0.2000, 0.2000, 0.2000])
---------end backward---------

对比上一例可以发现a的梯度从0.4被裁剪到了0.2，这里使用的clamp_是直接原地修改，所以不需要返回值。

也可将上述例子中的hook更改为有返回值的函数，效果相同。

部分例子参考：https://zhuanlan.zhihu.com/p/662760483

2. torch.nn.Module.register_forward_hook(hook_func)

除了register_hook是对tensor操作的hook之外，其他的hook都是对module进行操作的，这里的module包括各种layer，例如：Conv2d, Linear等

register_forward_hook在执行module的forward函数之后执行，其函数签名为

def hook(module, inputs, outpus):pass

注意：这里的module是当前被注册的module，inputs是执行forward之前的inputs，而outputs则是执行forward之后的outputs ，这么设计可能是为了方便读取执行之前的intputs。

如下例所示：

import torch
import torch.nn as nn# 定义一个简单的模块
class MyModule(nn.Module):def forward(self, x):print('forward'.center(20, '-'))return x * 2  # 假设这个模块简单地将输入乘以2# 创建模块实例
module = MyModule()# 定义一个hook函数，它接受输入和输出作为参数
def my_hook(module, input, output):print(f"Input: {input}")print(f"Output: {output}")# 注册hook函数
module.register_forward_hook(my_hook)# 创建一个输入张量
input_tensor = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)# 执行前向传播，这将触发hook函数的调用
output_tensor = module(input_tensor)

输出：

------forward-------
Input: (tensor([1., 2., 3.], requires_grad=True),)
Output: tensor([2., 4., 6.], grad_fn=<MulBackward0>)

从中我们可以看到，这里的Input还是执行forward之前的input，但是outputs是执行forward之后的outputs，从打印的------forward-------位置可以知道，这里的forward函数是在执行之后调用的hook。

我们可以使用hook实现torchsummary类似的功能，查看resnet18的各个层的输出情况，如下例

import torch
from torch import nn
from torchvision.models import resnet18class Visualize(nn.Module):def __init__(self, model) -> None:super().__init__()self.model = model# Register a hook for each layerfor name, layer in self.model.named_children():# add a property dynamicallylayer.name = name# module.name is the newly added propertylayer.register_forward_hook(lambda module, inputs, outputs:print(f"{module.name}".ljust(10), '-->', f'{outputs.shape}'))def forward(self, x):return self.model(x)model = resnet18()
inputs = torch.randn(1, 3, 224, 224)
vis = Visualize(model)
output = vis(inputs)

输出：

conv1      --> torch.Size([1, 64, 112, 112])
bn1        --> torch.Size([1, 64, 112, 112])
relu       --> torch.Size([1, 64, 112, 112])
maxpool    --> torch.Size([1, 64, 56, 56])
layer1     --> torch.Size([1, 64, 56, 56])
layer2     --> torch.Size([1, 128, 28, 28])
layer3     --> torch.Size([1, 256, 14, 14])
layer4     --> torch.Size([1, 512, 7, 7])
avgpool    --> torch.Size([1, 512, 1, 1])
fc         --> torch.Size([1, 1000])

如果使用使用apply对hook进行注册，apply会递归地将model里面的所有layer都进行相同的操作，于是结果就和for name, layer in self.model.named_modules()类似。

import torch
from torch import nn
from torchvision.models import resnet18def hook(module, inputs, outputs):print(module.__class__.__name__.ljust(10), end='')print(outputs.shape)def register(module):if isinstance(module, nn.Conv2d):module.register_forward_hook(hook)model = resnet18()
inputs = torch.randn(1, 3, 224, 224)
# 这里的apply会递归地把所有层都遍历，因此register_forward_hook注册到的层
# 是所有的Conv2d，包括子层，子层中的子层...
model.apply(register)
outputs = model(inputs)

输出为：

Conv2d    torch.Size([1, 64, 112, 112])
Conv2d    torch.Size([1, 64, 56, 56])
Conv2d    torch.Size([1, 64, 56, 56])
Conv2d    torch.Size([1, 64, 56, 56])
Conv2d    torch.Size([1, 64, 56, 56])
Conv2d    torch.Size([1, 128, 28, 28])
Conv2d    torch.Size([1, 128, 28, 28])
Conv2d    torch.Size([1, 128, 28, 28])
Conv2d    torch.Size([1, 128, 28, 28])
Conv2d    torch.Size([1, 128, 28, 28])
Conv2d    torch.Size([1, 256, 14, 14])
Conv2d    torch.Size([1, 256, 14, 14])
Conv2d    torch.Size([1, 256, 14, 14])
Conv2d    torch.Size([1, 256, 14, 14])
Conv2d    torch.Size([1, 256, 14, 14])
Conv2d    torch.Size([1, 512, 7, 7])
Conv2d    torch.Size([1, 512, 7, 7])
Conv2d    torch.Size([1, 512, 7, 7])
Conv2d    torch.Size([1, 512, 7, 7])
Conv2d    torch.Size([1, 512, 7, 7])

apply将所有的Conv2d都注册了，所以输出了所有的Conv2d的输出shape。

3.torch.nn.Module.register_backward_hook

在了解了前一个hook的用法之后，这个hook的作用也就不言而喻了，在backward之后执行，这里的hook函数签名如下

def hook_fn(module, grad_in, grad_out):pass

输入参数包括三个，分别是module，grad_in和grad_out，其中，grad_in和grad_out分别指代当前模块的输入和输出的梯度信息，若grad_in和grad_out包括多个输入输出，则grad_in和grad_out以元组形式呈现。

现在使用会register_backward_hook爆出警告：

module.py:1352: UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes "

解决办法就是使用新的hook函数register_full_backward_hook，新的hook函数功能更加强大，不仅仅包括模块的输入输出梯度信息，还包括内部的一些其他变量的梯度信息，但是register_backward_hook和register_full_backward_hook两者之间的兼容性并不是很完美。

示例

import torch
from torch import nn
from torchvision.models import resnet18def hook_fn(module, grad_in, grad_out):# 当前module的输入和输出梯度# 若module有多个输入，则grad_in为一个元组# y = wx+bprint(module.__class__.__name__)print("------------Input Grad------------")# 容错处理，部分元组中的变量会是Nonefor grad in grad_in:try:print(grad.shape)except AttributeError: print ("None found for Gradient")print("------------Output Grad------------")for grad in grad_out:  try:print(grad.shape)except AttributeError: print ("None found for Gradient")print("\n")net = resnet18()
for name, layer in net.named_children():# 每一个大的子层都注册一个勾子函数layer.register_backward_hook(hook_fn)# 为了能够执行backward，构建一些虚拟的输入输出
dummy_inputs = torch.randn(10, 3, 224, 224)
dummy_labels = torch.randint(0, 1001, (10, ))
loss_fn = nn.CrossEntropyLoss()y_hat = net(dummy_inputs)loss = loss_fn(y_hat, dummy_labels)
loss.backward()

输出：

module.py:1352: UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes "Linear
------------Input Grad------------
torch.Size([1000])
torch.Size([10, 512])
torch.Size([512, 1000])
------------Output Grad------------
torch.Size([10, 1000])AdaptiveAvgPool2d
------------Input Grad------------
torch.Size([10, 512, 7, 7])
------------Output Grad------------
torch.Size([10, 512, 1, 1])Sequential
------------Input Grad------------
torch.Size([10, 512, 7, 7])
------------Output Grad------------
torch.Size([10, 512, 7, 7])Sequential
------------Input Grad------------
torch.Size([10, 256, 14, 14])
------------Output Grad------------
torch.Size([10, 256, 14, 14])Sequential
------------Input Grad------------
torch.Size([10, 128, 28, 28])
------------Output Grad------------
torch.Size([10, 128, 28, 28])Sequential
------------Input Grad------------
torch.Size([10, 64, 56, 56])
------------Output Grad------------
torch.Size([10, 64, 56, 56])MaxPool2d
------------Input Grad------------
torch.Size([10, 64, 112, 112])
------------Output Grad------------
torch.Size([10, 64, 56, 56])ReLU
------------Input Grad------------
torch.Size([10, 64, 112, 112])
------------Output Grad------------
torch.Size([10, 64, 112, 112])BatchNorm2d
------------Input Grad------------
torch.Size([10, 64, 112, 112])
torch.Size([64])
torch.Size([64])
------------Output Grad------------
torch.Size([10, 64, 112, 112])Conv2d
------------Input Grad------------
None found for Gradient
torch.Size([64, 3, 7, 7])
None found for Gradient
------------Output Grad------------
torch.Size([10, 64, 112, 112])

最上面是警告信息可以忽略，然后根据backward的路径，从后往前进行返回。

使用如下代码查看resnet18的层级情况：

for name, layer in net.named_children():print(name)

输出：

conv1
bn1
relu
maxpool
layer1
layer2
layer3
layer4
avgpool
fc

可以看到这里的10个层对应上面hook函数返回的10个层。

综合以上两个部分，用一个示例演示同时构建前向和后向勾子函数：

import torch
import torch.nn as nn# 前向钩子示例
def forward_hook(module, input, output):print("{} forward hook:".format(module.__class__.__name__))print("Input:", input)print("Output:", output)print("")# 反向钩子示例
def backward_hook(module, grad_input, grad_output):print("{} backward hook:".format(module.__class__.__name__))print("Gradient input:")for item in grad_input:if item is not None:print(item.shape)print("Gradient output:")for item in grad_output:if item is not None:print(item.shape)print("")# 示例模型
class SimpleModel(nn.Module):def __init__(self):super(SimpleModel, self).__init__()self.fc1 = nn.Linear(10, 20)self.fc2 = nn.Linear(20, 1)def forward(self, x):x = torch.relu(self.fc1(x))x = self.fc2(x)return x# 示例
model = SimpleModel()# 注册前向钩子
hook_handle = model.fc1.register_forward_hook(forward_hook)# 注册反向钩子
hook_handle2 = model.fc1.register_backward_hook(backward_hook)# 示例输入数据
input_data = torch.randn(1, 10)# 前向传播
output = model(input_data)# 反向传播
loss = output.sum()
loss.backward()# 移除钩子
hook_handle.remove()
hook_handle2.remove()

输出：

Linear forward hook:
Input: (tensor([[-1.6549, -1.1471, -0.2341,  0.1456,  0.6528, -1.0562,  0.1078,  0.9752,0.8794,  1.0463]]),)
Output: tensor([[-0.6406,  0.0515,  0.1893, -0.5211, -0.2393,  0.2923,  0.0143,  0.6929,-0.4688, -0.1708, -0.6461,  0.5460, -0.1515, -0.1707, -0.5409, -0.6382,-0.9836,  0.3446,  0.2147, -0.7682]], grad_fn=<AddmmBackward0>)Linear backward hook:
Gradient input:
torch.Size([20])
torch.Size([10, 20])
Gradient output:
torch.Size([1, 20])

使用hook机制可视化resnet的特征图输出

import cv2
from torchvision import transforms
from torchvision.models import ResNet18_Weights, resnet18
import torch
import matplotlib.pyplot as pltdef viz(name):def imshow(module, input, output):feature_maps = input[0]# feature map dimension:# (batch_size, ch, width, height)# visualize 4 channels at mostmax_ch = min(feature_maps.size(1), 4)imgs = feature_maps[0, :max_ch, :, :]# print(imgs.shape)plt.figure(figsize=(12, 2))for i, img in enumerate(imgs):plt.subplot(1, 4, i+1)# plt.imshow(img.cpu(), cmap='gray')plt.imshow(img.cpu())plt.axis('off')if i == 0:plt.title(name)plt.show()return imshowdef main():trans = transforms.Compose([transforms.ToPILImage(),transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])])device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1).to(device)for name, module in model.named_modules():# 这里只对卷积层的feature map进行显示if isinstance(module, torch.nn.Conv2d):module.register_forward_hook(viz(name))img = cv2.imread(r'faces\ftw1.jpg')img = trans(img).unsqueeze(0).to(device)with torch.no_grad():model(img)main()

输出示例：

在这里插入图片描述
在这里插入图片描述

总结：

  1. 勾子函数可以在不修改源代码的情况下实现功能的注入
  2. 实现过程需要重写对应的勾子函数，需要注意执行的顺序以及参数的含义
    • register_forward_hook：在forward函数之后执行，输入参数为input和output，其中input是forward函数之前的输入，output是forwad函数之后的输入。这个勾子函数一般用于可视化特征图
    • register_backward_hook：在执行backward之时执行，backward到哪一个层就执行哪一个层的勾子函数，需要注意的是，输入参数分别为当前层的梯度输入和梯度输出，也即grad_input, grad_output，再者，使用该函数不能有原地修改的操作，否则会报异常。

参考内容

  • 一文搞懂PyTorch Hook
  • Pytorch官方文档
    PyTorch Hook用法解析

本文来自互联网用户投稿，该文观点仅代表作者本人，不代表本站立场。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。如若转载，请注明出处：http://www.hqwc.cn/news/576234.html

如若内容造成侵权/违法违规/事实不符，请联系编程知识网进行投诉反馈email:809451989@qq.com，一经查实，立即删除！

相关文章

人大金仓数据库介绍与使用指南

人大金仓数据库介绍与使用指南

人大金仓数据库是一款强大的关系型数据库管理系统&#xff0c;具有简单易用、高性能和稳定可靠的特点。本文将介绍人大金仓数据库的安装方法、常用的SQL语法以及相关工具的使用。 一、安装方法&#xff1a; 1、下载人大金仓数据库安装程序&#xff1b; 2、运行安装程序&#…
阅读更多...
康耐视visionpro-CogFindCircleTool工具详细说明

康耐视visionpro-CogFindCircleTool工具详细说明

CogFindCircleTool功能说明： 通过用多个卡尺找到多个点来拟合所要找的圆 CogFindCircleTool操作说明： ①.打开工具栏，双击或点击鼠标拖拽添加CogFindCircleTool工具 ②.添加输入图像，右键“链接到”或以连线拖拽的方式选择相应输入源 ③.预期的圆弧：设置预期圆弧的中心点…
阅读更多...
基于SSM+Jsp+Mysql的医院远程诊断系统

基于SSM+Jsp+Mysql的医院远程诊断系统

开发语言&#xff1a;Java框架&#xff1a;ssm技术&#xff1a;JSPJDK版本&#xff1a;JDK1.8服务器&#xff1a;tomcat7数据库&#xff1a;mysql 5.7&#xff08;一定要5.7版本&#xff09;数据库工具&#xff1a;Navicat11开发软件&#xff1a;eclipse/myeclipse/ideaMaven包…
阅读更多...
软考 系统架构设计师系列知识点之云原生架构设计理论与实践（8）

软考 系统架构设计师系列知识点之云原生架构设计理论与实践（8）

接前一篇文章&#xff1a;软考 系统架构设计师系列知识点之云原生架构设计理论与实践&#xff08;7&#xff09; 所属章节&#xff1a; 第14章. 云原生架构设计理论与实践 第2节 云原生架构内涵 14.2 云原生架构内涵 关于云原生的定义有众多版本&#xff0c;对于云原生架构的…
阅读更多...
arp 协议

arp 协议

数据链路层 我们之前学习到的 IP 协议解决的是数据跨网络传输的问题。 数据链路层解决的是&#xff1a;直接相连的主机&#xff0c;进行数据交付的问题&#xff01; 直接相连的设备包括我们的电脑&#xff0c;路由器等等哈&#xff01; 我们在网络基础那篇文章中讲过什么是以…
阅读更多...
指标监控和归因分析——数据异常波动

指标监控和归因分析——数据异常波动

目录 前言 一、基于统计分析检测指标异常 二、指标异常归因分析 2.1 横向归因分析 2.2 纵向归因分析 三、智能指标波动监控&归因分析 3.1 指标看板查看 3.2 指标归因分析 前言 企业搭建完善&#xff0c;全面的指标体系是企业数据指导业务经营决策的第一步&#xff…
阅读更多...
206基于matlab的无人机航迹规划(UAV track plannin）

206基于matlab的无人机航迹规划(UAV track plannin）

基于matlab的无人机航迹规划(UAV track plannin&#xff09;。输入输出参数包括 横滚、俯仰、航向角&#xff08;单位&#xff1a;度&#xff09;&#xff1b;横滚速率、俯仰速率、航向角速率&#xff08;单位&#xff1a;度/秒&#xff09;&#xff1b;飞机运动速度——X右翼、…
阅读更多...
文件上传失败原因分析与解决

文件上传失败原因分析与解决

图片文件上传失败 问题描述&#xff1a;在前端开发时&#xff0c;需要通过表单元素上传图片或其他文本&#xff0c;但是上传不成功&#xff0c;后端接口也没问题 html <!--onChange用来绑定数据 handleUpload用来提交数据--><form onSubmit{handleUpload}><…
阅读更多...
row_number 函数和关联更新

row_number 函数和关联更新

生成测试数据&#xff0c;房间号数据如下&#xff1a; CREATE TABLE hotel (floor_nbr,room_nbr) ASSELECT 1,100 FROM DUAL UNION ALLSELECT 1,100 FROM DUAL UNION ALLSELECT 2,100 FROM DUAL UNION ALLSELECT 2,100 FROM DUAL UNION ALLSELECT 3,100 FROM DUAL; 里面的房间号…
阅读更多...
redis-shake可视化监控

redis-shake可视化监控

目录 一.redis-shake v4 1.镜像 2.shake.toml 3.启动redis-shake后 二.json-exporter配置 1.Dockerfile 2.config.yml 三.prometheus配置 1.prometheus.yml 2.redis-shake.json 四.grafana 一.redis-shake v4 1.镜像 ######################### Dockerfile #########…
阅读更多...
StarRocks实战——多点大数据数仓构建

StarRocks实战——多点大数据数仓构建

目录 前言 一、背景介绍 二、原有架构的痛点 2.1 技术成本 2.2 开发成本 2.2.1 离线 T1 更新的分析场景 2.2.2 实时更新分析场景 2.2.3 固定维度分析场景 2.2.4 运维成本 三、选择StarRocks的原因 3.1 引擎收敛 3.2 “大宽表”模型替换 3.3 简化Lambda架构 3.4 模…
阅读更多...
计算机网络-RIP动态路由协议简介

计算机网络-RIP动态路由协议简介

一、概述 前面我们学习了动态路由协议按照工作机制及算法划分可以分为&#xff1a;距离矢量路由协议DV型和链路状态路由协议LS型。RIP就是典型的距离矢量路由协议&#xff0c;但是实际工作中用得已经比较少了。 距离矢量路由协议DV: RIP 链路状态路由协议LS: OSPF IS-IS 二、RI…
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023