mc_dropout_mode

tiatoolbox . models . engine . mc_dropout . mc_dropout_mode

class mc_dropout_mode(model)

Context manager: enable only Dropout layers, everything else stays in eval.

Example

Run 3 stochastic forward passes with dropout active while keeping BatchNorm in eval mode:

>>> model = nn.Module()
>>> images = '/path/to/images'
>>> model.eval()
>>> with mc_dropout_mode(model):
...     samples = [model(images) for _ in range(3)]
>>> # on exit, all modules are back to their previous training state

Initialize mc_dropout_mode.

Parameters:

model (nn.Module) – The model whose Dropout layers will be put in train mode inside the context.

Methods