Hi there,
I am implementing a bayesian neural network. For this I need the dropout layers to be active during inference while the rest of the layers are in eval mode. This is done with
model.eval()
model.classifier_layer[2].train(True) # Dropout layer
This works perfectly fine. However, I want to implement it on real-time on a raspberry pi. I export the model with torch.onnx.export(). For this I found the argument training of the onnx.export function. Here I have 3 choices: TrainingMode.EVAL, TrainingMode.TRAINING, TrainingMode.PRESERVE. I can’t use training or eval because then the whole model would be in training/eval mode. If I apply the TrainingMode.PRESERVE and export the model. The output of the model is always the same which implies that the dropout layer is not in training mode.
Is there any way I can get only the dropout layer to be in the training mode when exporting the model or am I doing anything wrong?
Many thanks, I appreciate the help!