Pytorch中的Hook（常用于网络特征可视化等）

最新推荐文章于 2025-06-02 16:01:20 发布

原创

最新推荐文章于 2025-06-02 16:01:20 发布 · 2.5k 阅读

35 ·

CC 4.0 BY-SA版权

文章标签：

#可视化 #深度学习

本文详细介绍了PyTorch中的Hook机制，包括如何利用Hook函数获取网络中的中间变量，如梯度值、输入输出值等，并通过具体示例展示了如何在实际应用中使用这些Hook函数。

Pytorch中的Hook

使用Hook函数获取网络中间变量.
Hook函数机制是不改变函数主体，实现额外功能，像一个挂件，挂钩。正是因为PyTorch计算图动态图的机制，所以才会有Hook函数。在动态图机制的运算，当运算结束后，一些中间变量就会被释放掉，例如，特征图，非leaf节点的梯度。但是有时候，我们需要这些中间变量，所以就出现了Hook函数。
torch提供了四种hook方法，分别用于

获取各个参数的梯度值。tensor.register_hook(hook)
获取各个层前向传播的输入输出值。Module.register_forward_hook(hook)
获取各个层前向传播的输入值。Module.register_forward_pre_hook(hook)
获取各个层反向传播的梯度值。Module.register_backward_hook(hook)

Torch.tensor.register_hook(hook)

记录中间各参数的梯度值

import torch

x=torch.tensor([3.],requires_grad=True)
y=torch.tensor([5.],requires_grad=True)
print('x:',x)
print('y:',y)

a=x+y # a=3+5=8
b=x*y # b=3*5=15

c=a*b # c=8*15=120

# 存储a的梯度
a_grads=[]
def hook_grad(grad):
    a_grads.append(grad)

handle=a.register_hook(hook_grad)
c.backward()

x: tensor([3.], requires_grad=True)
y: tensor([5.], requires_grad=True)

# c=x^2*y+x*y^2。c对x求导：2xy+y^2,c对y求导:x^2+2xy。x=3，y=5
print(x.grad,y.grad,a.grad,b.grad,c.grad)

tensor([55.]) tensor([39.]) None None None

# c=a*x*y,对a求导为：xy.(x=3,y=5)
print(a_grads)
# c=(x+y)*b,对b求导为：x+y
print(b.grad)

[tensor([15.])]
None

# 移除hook函数
handle.remove()

torch.nn.Module.register_forward_hook

注册module的前向传播Hook函数
记录前向传播的中生成的中间输入输出结果。

module:当前网络层
input:本层输入数据。与register_forward_hook中得到的输入数据相同
output:本层输出数据

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) # in_c,out_c,f_size
        self.pool1=nn.MaxPool2d(2)
    def forward(self,x):
        x=self.conv1(x)
        x=self.pool1(x)
        return x