How to delete some layers in pretrained model?

There are several options to remove/replace this layer depending on your use case.

You could copy the source code of the model, remove the layer, change the forward method, and store this new model definition as a custom model. This would allow you to load this new model in your scripts and just train it if needed. The state_dict would not contain this layer (and in particular its parameters) anymore and you would not be able to directly load a pretrained state_dict anymore (using strict=False as a workaround should work).

Alternatively, you could also load the original model and replace the unwanted layer with an nn.Identity. While this approach might be simpler, you would need to replace this layer after each model creation.

3 Likes