Change output shape of any CNN architecture

Is there a general way to change the output shape of any Pytorch architecture?

Based on this article (Finetuning Torchvision Models — PyTorch Tutorials 1.2.0 documentation) it seems like each architecture requires a different method to change the number of output nodes (classes) - for instance,
for resnet18 you can replace model.fc with a nn.Linear layer but for vgg11_bn you have to replace the model.classifier[6] layer. Unfortunately, that article does not have all of the currently implemented PyTorch architectures. In general as more architectures are added it would be hard to keep adding new wrapper functions for changing the output shape. It would be very helpful if there were a consistent way to modify the output shape of any pytorch architecture.

1 Like