CODEWITHSUNDEEP
pythonpytorch
Siladittya
Siladittya

Reputation: 1205

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

I am trying to extract features from an intermediate of ResNet18 in PyTorch using forward hooks

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)

I also want the predictions alongside the feature outputs from layer 4

But I am getting the TypeError: hook() takes 3 positional arguments but 4 were given

The full error message is this

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

Why is the hook not working, although in various forums I see that this is the way to do it?

Upvotes: 0

Views: 1947

Answers (1)

Ivan
Ivan

Reputation: 40658

Here's a simple example of a forward hook, it must have three parameters model, input, and output:

m = models.resnet18(pretrained=False)

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

m.fc.register_forward_hook(hook)

Try it with dummy data:

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

To combine it with your nn.Module, you need to implement hook with an extra argument 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

Note - self corresponds to the CCLModel instance while model is the layer we're hooked on i.e. nn.Linear

Here's an example:

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

Upvotes: 1

Related Questions