How to delete layer in pretrained model?

2

Don’t you need this layer at all, i.e. do you want to get the avgpool back from your model?
If so, you could createn a module returning its input:

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


model = models.resnet18(pretrained=False)
model.fc = Identity()
x = torch.randn(1, 3, 224, 224)
output = model(x)
print(output.shape)

From the code it looks like you are using a resnet, so I used it in my examples.

However, usually you would like to use another linear layer with your number of classes as its output.

model.fc = nn.Linear(512, num_classes)
52 Likes