TypeError: hook() 需要 3 个位置 arguments 但给出了 4 个

Question

我正在尝试使用前向钩子从 PyTorch 中的 ResNet18 中间体中提取特征

class CCLModel(nn.Module):
    def __init__(self,output_layer,*args):
        self.output_layer = output_layer
        super().__init__(*args)

        self.output_layer = output_layer
        #PRETRAINED MODEL
        self.pretrained = models.resnet18(pretrained=True)
    
        #TAKING OUTPUT FROM AN INTERMEDIATE LAYER

        #self._layers = []
        for l in list(self.pretrained._modules.keys()):
            #self._layers.append(l)
            if l == self.output_layer:
                handle = getattr(self.pretrained,l).register_forward_hook(self.hook)
   
    def hook(self,input,output):
        return output

    def _forward_impl(self, x):
        x = self.pretrained(x)
        return x

    def forward(self, x):
        return self._forward_impl(x)

我还想要第 4 层的特征输出旁边的预测

但是我收到了TypeError: hook() takes 3 positional arguments but 4 were given

完整的错误信息是这样的

TypeError                                 Traceback (most recent call last)
<ipython-input-66-18c4a0f917f2> in <module>()
----> 1 out = model(x.to('cuda:0').float())

6 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

<ipython-input-61-71fe0d1420a6> in forward(self, x)
     78 
     79     def forward(self, x):
---> 80         return self._forward_impl(x)
     81 
     82     '''def forward(self,x):

<ipython-input-61-71fe0d1420a6> in _forward_impl(self, x)
     73         #x = torch.flatten(x, 1)
     74         #x = self.fc(x)
---> 75         x = self.pretrained(x)
     76 
     77         return x

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.6/dist-packages/torchvision/models/resnet.py in forward(self, x)
    218 
    219     def forward(self, x):
--> 220         return self._forward_impl(x)
    221 
    222 

/usr/local/lib/python3.6/dist-packages/torchvision/models/resnet.py in _forward_impl(self, x)
    209         x = self.layer2(x)
    210         x = self.layer3(x)
--> 211         x = self.layer4(x)
    212 
    213         x = self.avgpool(x)

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    729                 _global_forward_hooks.values(),
    730                 self._forward_hooks.values()):
--> 731             hook_result = hook(self, input, result)
    732             if hook_result is not None:
    733                 result = hook_result

TypeError: hook() takes 3 positional arguments but 4 were given

为什么钩子不起作用，尽管在各种论坛中我看到这是这样做的方法？

Answer 1

这是一个前向挂钩的简单示例，它必须具有三个参数model 、 input和output ：

m = models.resnet18(pretrained=False)

def hook(module, input, output):
    print(output.detach().shape)

m.fc.register_forward_hook(hook)

尝试使用虚拟数据：

>>> m(torch.rand(1, 3, 224, 224))
torch.Size([1, 1000])
<<< tensor(...)

要将它与您的nn.Module结合起来，您需要使用额外的参数self实现钩子：

class CCLModel(nn.Module):
    def __init__(self, output_layer, *args):
        super(CCLModel, self).__init__()

        self.pretrained = models.resnet18(pretrained=True)
        self.output_layer = output_layer

        self.output_layer.register_forward_hook(self.hook)
   
    def hook(self, module, input, output):
        return print(output.shape)

    def forward(self, x):
        x = self.pretrained(x)
        x = self.output_layer(x)
        return x

注意- self对应于CCLModel实例，而model是我们迷上的层，即nn.Linear

这是一个例子：

>>> m = CCLModel(nn.Linear(1000, 100))    
>>> m(torch.rand(1, 3, 224, 224))
torch.Size([1, 100])
<<< tensor(...)