Separate Torchscript model into two parts

Hello, all.

Please imagine that we want to separate the process of the Deep learning model into the feature extraction process and the classification process.
For example, the feature extraction process of AlexNet is from the first convolution layer and the last pooling layer which is located in front of the first fully connected layer. The remaining part from the first fully connected layer to the end of the Deep learning model is the classification process.

I checked that this partitioning can be achieved for the pre-trained AlexNet model in the Torchvision package. The feature extraction process is implemented in features which is a nn.Sequential instance and the classification process is implemented in classifier which is also a nn.Sequential instance.
Thus, model.features.forward(input data) gives me the result of the feature extraction and model.classifier.forward(feature data) gives me the result of the classification.

My question is whether this is also possible for the Torchscript model.
According to the tutorial of Torchscript, we use model, filename) to convert a PyTorch model to a Torchscript model.
Then, when I load the saved Torchscript model, it is impossible to access features directly like the case of PyTorch model.

The reason what I thought is that the PyTorch model is eager mode and the Torchscript model is graph mode. Thus, the graph mode model might not able to be separated.

  • If my idea is correct, is there no way to separate graph mode model into two parts without conducting for features and classifier?
  • If my idea is not correct, please share your idea.

Many thanks !


You can export custom methods with TorchScript. So, I would make your model to have a “extract_features” method that extracts them y mark it as @jit.export . So, you can call model.extract_features((x).