NotImplemented Error after dropping layer on pretrained network

Hi All. I have a trained Xception model that I want to use as a backbone in another network. In particular I want to drop all the layers starting with the global average pooling. In the code below this means dropping the last block (GAPClassifier).

The problem is that after doing so I get a NotImplementedError.

I’m assuming that the issue is how I’m dropping the last block, namely

model=torch.nn.Sequential(*(list(model.children())[:-1]))

Are there situations where dropping a block is more complicated? Whats the best way to handle this?

Here is a run through of what is happening (unfortunately I can’t include the entire model in this post). Thanks!

test_batch=torch.rand((3,14,SIZE,SIZE)).cuda()
model(test_batch) #<== executes successfully

def block_names(model,end=None,start=None):
    layer_list=list(model.children())
    for i,layer in enumerate(layer_list[start:end]):
        print(i,layer._get_name())

block_names(model)
"""OUTPUT
0 EntryBlock
1 ModuleList
2 SeparableStack
3 XBlock
4 SeparableStack
5 GAPClassifier
"""

model=torch.nn.Sequential(*(list(model.children())[:-1]))

block_names(model)
"""OUTPUT
0 EntryBlock
1 ModuleList
2 SeparableStack
3 XBlock
4 SeparableStack
"""
model(test_batch) # throws NotImplementedError

By wrapping submodules in an nn.Sequential block you are assuming all children are applied sequentially in the original model.
While this might be often the case, some models are a bit more complicated so that this simple container won’t restore the original work flow anymore or will throw an error.
E.g. this .view() operation will be missing, if you are extracting all child modules from this model.

I’m not sure, which module throws the NotImplementedError and it’s a bit strange, as I would expect some shape mismatch etc.

If you just want to remove the last layer, you could replace it with an nn.Identity module.

That makes sense. Thanks!