torch.nn.Module.register_forward_hook使用

本文简单介绍 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)