本文简单介绍 torch.nn.Module.register_forward_pre_hook钩子函数的使用,简单写了一个卷积的网络,在net.conv1.register_forward_pre_hook注册钩子函数,函数参数为module与输入input数据,在self.conv1层注册,模型先运行forward_pre_hook函数,参数为module与input(你可以通过forward_pre_hook修改模块)后执行x = self.conv1(input)。
torch.nn.Module.register_forward_hook文章:https://www.cnblogs.com/tangjunjun/p/17478085.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, input):
x = self.conv1(input)
return x
def forward_pre_hook(module, data_input):
input_block.append(data_input)
module.weight.data=torch.ones((module.weight.shape)) # 更改权重
net = Net()
input_block = list()
handle = net.conv1.register_forward_pre_hook(forward_pre_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)
修改权重运行结果:
我将权重设定固定值,且关掉bias,则y=wx,因此为固定值
不修改权重运行结果:
因为self.conv1将默认随机权重,因此值不一样。