Quick example to show how to activate and deactivate mcdropout

mc_dropout_context.py

import numpy as np
import torch
from torchvision.models import vgg16

from baal.bayesian.dropout import MCDropoutModule
from baal.modelwrapper import ModelWrapper

model = vgg16()
wrapper = ModelWrapper(model, None)
input = torch.randn([2, 3, 64, 64])


def is_deterministic(preds):
    return all([np.allclose(preds[..., i], preds[..., 0], rtol=1e-3) for i in range(1, preds.shape[-1])])


with MCDropoutModule(model):
    model.eval()
    # The model will not be deterministic
    assert not is_deterministic(wrapper.predict_on_batch(input, iterations=10))
# Is deterministic again.
model.eval()
assert is_deterministic(wrapper.predict_on_batch(input, iterations=10))