本文简单介绍 torch.nn.Module.register_forward_hook钩子函数的使用,简单写了一个卷积的网络,在net.conv1.register_forward_hook注册钩子函数,则会有module、输入input数据与卷积后输出数据output,重点说明module是关于模型结构self.conv1模块,在self.conv1层注册,模型先运行x = self.conv1(input)该层卷积,后执行forward_hook(module, input,output)该函数,输入为module、input与output,此时修改self.conv1权重等都不影响x = self.conv1(input)执行结果。
torch.nn.Module.register_forward_pre_hook文章:https://www.cnblogs.com/tangjunjun/p/17477796.html
代码:
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,bias=False)
def forward(self, x):
x = self.conv1(x)
return x
def forward_hook(module, input,output):
'''
Args:
module: 模型模块
input: 输入
output: 输出
'''
input_block.append(input)
output_block.append(output)
# module.weight.data=torch.ones((module.weight.shape)) # 更改权重
net = Net()
input_block = list()
output_block=list()
handle = net.conv1.register_forward_hook(forward_hook) # 在conv1中注册
if __name__ == '__main__':
# inference
fake_img = torch.ones((1, 1, 4, 4)) # batch size * channel * H * W
output = net(fake_img)
print(output)