모델의 특정 계층의 입력과 출력을 정확하게 찾으려면 Register_forward_hook()을 사용하십시오.

어떤 방법에서는 모델의 마지막 레이어의 성능을 분류기로 비교하기 위해 모델 중간에 있는 히든 레이어를 분류자로 사용한다는 내용을 우연히 논문에서 읽었기 때문에 어떻게 하면 쉽고 빠르게 출력을 뽑아낼 수 있을지 고민했습니다. 모델의 특정 레이어를 꺼낸 후 이를 수행할 수 있는 기성 후크 기능이 있음을 발견했습니다.

훅

Hook은 네트워크 내 특정 계층의 입력, 출력, 기타 정보를 Hooking하는 데 사용되는 Hook입니다. 네트워크 내 특정 계층의 자세한 정보를 알고 싶다면 별도로 작성할 필요가 없습니다. 네트워크를 정의할 때 인쇄하려면 Hook 함수를 직접 작성하면 됩니다.

Register_forward_hook

소스코드에 보면 Hook는 전달() 함수가 실행되기 전에만 사용할 수 있다고 명시되어 있고, 전달 함수가 실행된 후에 작성하면 소용이 없습니다. 즉, Hook을 실행하려면 먼저 Hook을 작성해야 한다는 뜻입니다. 기능을 수행한 다음 네트워크를 인스턴스화합니다.

def register_forward_hook(self, hook):
        r'''Registers a forward hook on the module.
        The hook will be called every time after :func:`forward` has computed an output.
        It should have the following signature::
            hook(module, input, output) -> None or modified output
        The hook can modify the output. It can modify the input inplace but
        it will not have effect on forward since this is called after
        :func:`forward` is called.

        Returns:
            :class:`torch.utils.hooks.RemovableHandle`:
                a handle that can be used to remove the added hook by calling
                ``handle.remove()``
       '''
        handle = hooks.RemovableHandle(self._forward_hooks)
        self._forward_hooks[handle.id] = hook
        return handle

질문

모델에 여러 개의 Linear 레이어가 있는 경우도 있지만 net.children()에서 추출한 동일한 유형의 모듈도 모두 동일한 이름을 가지므로 현재 Linear 레이어의 입력 및 출력 차원을 기준으로 판단하여 정확하게 이 레이어에 잠겨 있습니다. 다른 모듈도 여전히 적용 가능합니다.

코드 부분

import torch
import torch.nn as nn
class TestForHook(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_1=nn.Linear(2,2)
        self.linear_2 = nn.Linear(in_features=2, out_features=1)
        self.add_module('linear1', self.linear_1)
        self.add_module('linear2', self.linear_2)
        self.relu = nn.ReLU()
        self.relu6 = nn.ReLU6()
        self.initialize()
    def forward(self,x):
        linear_1=self.linear_1(x)
        linear_2=self.linear_2(linear_1)
        relu=self.relu(linear_2)
        relu_6 = self.relu6(relu)
        layers_in=(x,linear_1,linear_2)
        layers_out=(linear_1,linear_2,relu)
        return relu_6,layers_in,layers_out
    def initialize(self):
        '''定义特殊的初始化，用于验证hook作用时是否获取了权重'''
        self.linear_1.weight=torch.nn.Parameter(torch.FloatTensor([[1,1],[1,1]]))
        self.linear_1.bias=torch.nn.Parameter(torch.FloatTensor([1,1]))
        self.linear_2.weight=torch.nn.Parameter(torch.FloatTensor([[1,1]]))
        self.linear_2.bias=torch.nn.Parameter(torch.FloatTensor([1]))
        return True
#定义hook函数用来决定勾出来的网络信息用来做什么
#定义用于获取网络各层输入输出的tensor容器
#定义nodule_name用于记录相应的module名字
module_name=[]
features_in_hook=[]
features_out_hook=[]
#hook函数需要3个参数，这三个参数是系统传给hook函数的，自己不能修改这三个参数
#hook函数负责将获取的输入输出添加到feature列表中 并提供相应的module名字
def hook(module,input,output):
    print("hooker working")
    module_name.append(module.__class__)
    features_in_hook.append(input)
    features_out_hook.append(output)
    return None
#对需要的层register hook
#register hook必须在forward（）函数被执行之前，也就是实例化网络之前，下面的代码对网络除了ReLU以外的层都register了
#也可以选定其中的某些层进行register
net=TestForHook()
net_children=net.children()
#不同Linear层的参数in_features和out_features通常不同,可以用这些信息来判断
for child in net_children:
    if isinstance(child, nn.Linear) and child.in_features == 2 and child.out_features == 2:
    # if isinstance(child, nn.Linear):
        child.register_forward_hook(hook=hook)
#测试forward()提供的输入输出特征
x = torch.FloatTensor([[0.1, 0.1], [0.1, 0.1]])
out, features_in_forward, features_out_forward = net(x)
# print("*"*5+"forward return features"+"*"*5)
# print(features_in_forward)
# print(features_out_forward)
# print("*"*5+"forward return features"+"*"*5)
#hook通过list结构进行记录，所以可以直接print
print("*"*5+"hook record features"+"*"*5)
print(features_in_hook)
print(features_out_hook)
print(module_name)
print("*"*5+"hook record features"+"*"*5)