How to remove the adaptive average pool layer from vgg19?

I have loaded the pretrained model of vgg19. How to remove the adaptive average pool layer which is present before the classifier?

import torch.nn as nn
from torchvision.models import vgg19

net = vgg19(pretrained=True)
net.avgpool = nn.Identity()

But if you remove the average pooling layer, your output before classifier must have shape (N, 512, 7, 7), or you have to change the in_features of net.classifier[0] (the Linear layer).

Hi @Eta_C thanks for your reply. I am currently using pytorch 0.4.0 for my task so Identity function is not supported in this version i am supposing because of which it is unavailable. Is there any other alternative?

do you mean something like this,

net2 = nn.Sequential()
net2.add_module('features', nn.Sequential(*list(net.named_children())[0][1]))
net2.add_module('classifier', nn.Sequential(*list(net.named_children())[2][1]))

If you do not have nn.Identity, just define it.

import torch.nn as nn
from torchvision.models import vgg19

class Identity(nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()
    def forward(self, x):
        return x

net = vgg19(pretrained=True)
net.avgpool = Identity()

@Eta_C thanks for your solution :slight_smile:

I don’t know if in earlier versions of PyTorch the following works, but in v1.6 deleting a layer is as simple as:

del net.avgpool

This both removes the layer from model.modules and model.state_dict.
This is also does not create zombie layers, as an Identity layer would do; Simplifying model loading.

Double post from here with a warning, that this approach might break your model if the forward method isn’t fixed afterwards.