pytorch中的hook

最新推荐文章于 2025-02-19 20:58:29 发布

cccler

最新推荐文章于 2025-02-19 20:58:29 发布

阅读量433

点赞数

分类专栏： pytorch

本文链接：https://blog.csdn.net/qq_38992914/article/details/103122520

版权

本文介绍了PyTorch中的hook机制，包括观察tensor梯度、层的输入输出以及梯度。详细阐述了不同类型的hook函数及其应用场景，如forward_hook、backard_hook和tensor.register_hook。并提醒在使用后需要注意释放hook以避免内存消耗过大。还提供了简单的测试代码用于理解和应用hook。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

pytorch中的hook

为了方便，我都是用单输入单输出的结构来测试的。

pytorch提供三种hook，一种用来看tensor的grad，一种用来看一个层的输入输出，还有一种用来看一个层输入输出的梯度。

所有的hook函数，都需要定义一个function供它使用。

tensor.register_hook(function)是最好理解的，能够获取的就是这个tensor的梯度。它需要的function格式为hook(grad)，这里的grad就是tensor的梯度。

Conv2d.forward_hook(function)（当然其他一些layer也可以，这里就以Conv2d举例了），它需要的function格式为forward_hook(module,input,output)。这里的module就是这个Conv2d，input和output分别是前向传播时这个Conv2d的输入和输出。

Conv2d.backard_hook(function)，它需要的function格式为backward_hook(module,input_grad,output_grad)。这里的input和output，是反向传播时来说的（也就是和forward时相反），下面说的输入和输出，也是按反向传播的方向来说的。其中input_grad里有三个grad，按顺序分别是输入tensor的grad，layer自身weight的grad和layer自身bias的grad；output_grad中只有一个grad，就是输出tensor的grad。

这些hook，如果只是要打印一下需要的参数，比较方便，如果想要存起来供其他工具使用（比如tensorboardX），那么最好弄一个list来存储它们。

按照网络正常执行的顺序，forward_hook会最先输出，然后按照反向传播的顺序，backward_hook和tensor上的hook会依次输出。

不需要使用hook时一定要释放，否则显存会炸掉的…

测试代码：

import torch
import torch.nn as nn
import numpy as np

conv1=nn.Conv2d(1,1,3,1,1,bias=False)
conv2=nn.Conv2d(1,1,3,1,1,bias=False)
conv1.weight.data=torch.tensor([[[[1,2

最低0.47元/天解锁文章