Put a new top on ResNet

I’m trying to replace the last layer of an imagenet trained model.

model_imagenet = models.resnet18(pretrained=True)
model_imagenet

class ResNetNewTop(nn.Module):
    def __init__(self, old_model, num_classes=2):
        super(ResNetNewTop, self).__init__()
        self.bottom = nn.Sequential(*list(old_model.children())[:-1])
        self.top = nn.Linear(512, num_classes)  # TODO: input size (512) from bottom

    def forward(self, x):
        x = self.bottom(x)
        x = self.top(x)  # <<< Errors here
        return x

model = ResNetNewTop(model_imagenet, num_classes=2)
model.cuda()
model

This is throwing

RuntimeError: matrix and matrix expected at [snip] /torch/lib/THC/generic/THCTensorMathBlas.cu:235

Can I put Sequential() into a new model directly? Or do I have to loop through the layers and build up the model that way?

Yes you can do that. You likely forgot to merge some dimensions before giving them to the Linear layer (try print(x.size()) before self.top(x) - it will be 4D). Just do self.top(x.view(-1, 512)).

Additionally, in case of ResNet, you could just replace the final layer with a new one and it should work I think:

model_imagenet.fc = nn.Linear(512, num_classes)
2 Likes

I do this:

model = models.__dict__[opt.arch](pretrained=True) #pretrained=False if you don't want to use pre-trained weights.
model.fc = nn.Linear(512, num_classes)

I think it will only work for ResNets. Other models don’t use self.fc.

Yes. It will only work for ResNets. Your answer is more generalized.

Thank you! The diagnosis was spot on. I like the shortcut too; it helped me realize how loosely coupled all of the layers are.

So what is the right way to add more than one layers instead of just replacing the existing one? And i also want to keep the original name of each predefined layer.
In init, if I use

feat = nn.Sequential(*list(resnet.children())[:-1])
fc1 = nn.Linear(a, b)
fc1 = nn.Linear(b, c)

the layer names in feat are lost.