Inference on Specific Layers of a Model


I’m building a flexible module that implement Monte Carlo Dropout for Bayesian Inference on any kind of model.
(Implementing MC dropout by adding dropout to fc layers during test time and inference for multiple times to get distribution of outputs)
To save time, I want the to divide the model into two parts:

  1. No dropout implemented so run only once
  2. Layers which are after the first dropout (first fc layer) therefore will run multiple times.

I can’t use nn.Sequential() because many of the models I work with have functionality in forward (like x.permute()).

Do you have any Idea how to inference on only specific layers without changing the model itself or the forward() code?

Unfortunately not and I think the cleanest way would be to split the model into different parts and call them separately.

And how could I do that?
I am trying inference on the penultimate layer of an Alexnet Model but I couldn’t get it.

The author of this question was working on a custom model so could split the model into e.g. blocks which would fit their use case.
If you are working with a pretrained model created in e.g. torchvision, you could e.g. replace the last linear layer with an nn.Identity(), could write a custom model by deriving from AlexNet as its base, or use some helper methods from torchvision.