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])

相关推荐
张较瘦_几秒前
[论文阅读] 人工智能 + 软件工程 | 大模型破局跨平台测试!LLMRR让iOS/安卓/鸿蒙脚本无缝迁移
论文阅读·人工智能·ios
Pocker_Spades_A12 分钟前
Python快速入门专业版(十五):数据类型实战:用户信息录入程序(整合变量、输入与类型转换)
数据库·python
IMER SIMPLE16 分钟前
人工智能-python-深度学习-神经网络-GoogLeNet
人工智能·python·深度学习
钮钴禄·爱因斯晨19 分钟前
深入剖析LLM:从原理到应用与挑战
开发语言·人工智能
強云21 分钟前
cuda + cudnn + pytorch (win10)
pytorch
InternLM23 分钟前
专为“超大模型而生”,新一代训练引擎 XTuner V1 开源!
人工智能·开源·xtuner·书生大模型·大模型训练框架·大模型预训练·大模型后训练
小宁爱Python30 分钟前
Django 从环境搭建到第一个项目
后端·python·django
JT85839641 分钟前
AI GEO 优化能否快速提升网站在搜索引擎的排名?
人工智能·搜索引擎
幂律智能43 分钟前
吾律——让普惠法律服务走进生活
人工智能·经验分享
IT_陈寒1 小时前
Java性能优化:从这8个关键指标开始,让你的应用提速50%
前端·人工智能·后端