torch.nn.Module.register_forward_pre

本文简单介绍 torch.nn.Module.register_forward_pre_hook钩子函数的使用，简单写了一个卷积的网络，在net.conv1.register_forward_pre_hook注册钩子函数，则会有module与输入input数据，重点说明module是关于模型结构self.conv1模块，在self.conv1层注册，模型先运行forward_pre_hook函数，输入为module与input，此时你可以在forward_pre_hook函数中修改，修改完后才会执行x = self.conv1(input)此句代码。

代码：

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将默认随机权重，因此值不一样。

JZTXT

torch.nn.Module.register_forward_pre_hook使用