Resnet last layer modification

Currently you are rewrapping your pretrained resnet into a new nn.Sequential module, which will lose the forward definition. As you can see in this line of code in the original resnet implementation, the activation x will be flattened before being passed to the last linear layer. Since this is missing now, you’ll get the size mismatch error.

You could just manipulate the model.fc attribute and add your dropout and linear layer:

model = models.resnet18(pretrained=True)
model.fc = nn.Sequential(
    nn.Dropout(0.5),
    nn.Linear(num_ftrs, 10)
)

This approach will keep your forward method.

It depends on the criterion you are using if you should add a final non-linearity in your model.
If you are dealing with a classification use case and would like to use nn.CrossEntropyLoss, you should pass the raw logits to the loss function (i.e. no final non-linearity), since nn.LogSoftmax and nn.NLLLoss will be called internally.
However, if you would like to use nn.NLLLoss, you should add the nn.LogSoftmax manually.

16 Likes