torch.nn.Module.register_forward_pre_hook使用

本文简单介绍 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将默认随机权重，因此值不一样。