Onnx Export with Dropout Layer = Training, rest eval

1

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!

2

I’m not familiar with ONNX and don’t know why the PRESERVE mode is not working.
However, if for some reason the dropout layers are disabled during the export (have you verified it by e.g. just exporting a single dropout layer?), you could most likely use a custom module and hardcode the training argument:

class MyDropout(nn.Module):
    def __init__(self, p):
        self.p = p

    def forward(self, x):
        x = F.dropout(x, p=p, training=True)
        return x

and replace all layers with this custom implementation.