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()
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.