Pytorch--Hooks For Module

文章目录


1.register_module_forward_pre_hook

在 PyTorch 中,register_module_forward_pre_hook 是一个方法,用于向模型的模块注册前向传播预钩子(forward pre-hook)。预钩子是在模块的前向传播之前被调用的函数,允许在模块接收输入之前对输入进行修改或记录。

c 复制代码
import torch
import torch.nn as nn

# 定义一个前向传播预钩子函数
def forward_pre_hook(module, input):
    print("Forward pre-hook called for module:", module)
    print("Input shape:", input[0].shape)

# 创建一个模型类
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)

# 创建模型实例
model = MyModel()

# 注册前向传播预钩子
model.register_module_forward_pre_hook(forward_pre_hook)

# 输入数据
input_data = torch.randn(1, 10)

# 前向传播
output = model(input_data)
python 复制代码
Forward pre-hook called for module: Linear(in_features=10, out_features=10, bias=True)
Input shape: torch.Size([1, 10])

2.register_module_forward_hook

在 PyTorch 中,register_module_forward_hook 是一个方法,用于向模型的模块注册前向传播钩子(forward hook)。钩子是在模块的前向传播过程中被调用的函数,可以用于获取中间特征、对特征进行修改或记录等操作。

python 复制代码
import torch
import torch.nn as nn

# 定义一个前向传播钩子函数
def forward_hook(module, input, output):
    print("Forward hook called for module:", module)
    print("Input shape:", input[0].shape)
    print("Output shape:", output.shape)

# 创建一个模型类
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)

# 创建模型实例
model = MyModel()

# 注册前向传播钩子
model.register_forward_hook(forward_hook)

# 输入数据
input_data = torch.randn(1, 10)

# 前向传播
output = model(input_data)
python 复制代码
Forward hook called for module: Linear(in_features=10, out_features=10, bias=True)
Input shape: torch.Size([1, 10])
Output shape: torch.Size([1, 10])

3.register_module_backward_hook

在 PyTorch 中,register_module_backward_hook 是一个方法,用于向模型的模块注册反向传播钩子(backward hook)。钩子是在模块的反向传播过程中被调用的函数,可以用于获取梯度、对梯度进行修改或记录等操作。

python 复制代码
import torch
import torch.nn as nn

# 定义一个反向传播钩子函数
def backward_hook(module, grad_input, grad_output):
    print("Backward hook called for module:", module)
    print("Grad input shape:", grad_input[0].shape)
    print("Grad output shape:", grad_output[0].shape)

# 创建一个模型类
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)

# 创建模型实例
model = MyModel()

# 注册反向传播钩子
model.register_backward_hook(backward_hook)

# 输入数据
input_data = torch.randn(1, 10)
target = torch.randn(1, 10)

# 前向传播和反向传播
output = model(input_data)
loss = nn.MSELoss()(output, target)
loss.backward()
python 复制代码
Backward hook called for module: Linear(in_features=10, out_features=10, bias=True)
Grad input shape: torch.Size([1, 10])
Grad output shape: torch.Size([1, 10])

相关推荐
databook1 天前
Manim实现脉冲闪烁特效
后端·python·动效
程序设计实验室1 天前
2025年了,在 Django 之外,Python Web 框架还能怎么选?
python
飞哥数智坊1 天前
GPT-5-Codex 发布,Codex 正在取代 Claude
人工智能·ai编程
倔强青铜三1 天前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
虫无涯1 天前
Dify Agent + AntV 实战:从 0 到 1 打造数据可视化解决方案
人工智能
Dm_dotnet1 天前
公益站Agent Router注册送200刀额度竟然是真的
人工智能
算家计算1 天前
7B参数拿下30个世界第一!Hunyuan-MT-7B本地部署教程:腾讯混元开源业界首个翻译集成模型
人工智能·开源
用户2519162427111 天前
Python之语言特点
python
机器之心1 天前
LLM开源2.0大洗牌:60个出局,39个上桌,AI Coding疯魔,TensorFlow已死
人工智能·openai
刘立军1 天前
使用pyHugeGraph查询HugeGraph图数据
python·graphql