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

相关推荐
麻雀无能为力26 分钟前
CAU数据挖掘实验 表分析数据插件
人工智能·数据挖掘·中国农业大学
时序之心29 分钟前
时空数据挖掘五大革新方向详解篇!
人工智能·数据挖掘·论文·时间序列
IMPYLH1 小时前
Python 的内置函数 reversed
笔记·python
.30-06Springfield1 小时前
人工智能概念之七:集成学习思想(Bagging、Boosting、Stacking)
人工智能·算法·机器学习·集成学习
说私域2 小时前
基于开源AI智能名片链动2+1模式S2B2C商城小程序的超级文化符号构建路径研究
人工智能·小程序·开源
永洪科技2 小时前
永洪科技荣获商业智能品牌影响力奖,全力打造”AI+决策”引擎
大数据·人工智能·科技·数据分析·数据可视化·bi
shangyingying_12 小时前
关于小波降噪、小波增强、小波去雾的原理区分
人工智能·深度学习·计算机视觉
小赖同学啊3 小时前
物联网数据安全区块链服务
开发语言·python·区块链
码荼3 小时前
学习开发之hashmap
java·python·学习·哈希算法·个人开发·小白学开发·不花钱不花时间crud
书玮嘎3 小时前
【WIP】【VLA&VLM——InternVL系列】
人工智能·深度学习