sage.architecture.network.mc_dropout

Monte-Carlo dropout inference utilities.

At test time, keeping the dropout layers stochastic and averaging several forward passes turns the network into an approximate Bayesian model: the spread across passes is an estimate of epistemic uncertainty. With the heteroscedastic / consistency heads this complements the predicted (aleatoric) sigma — together they give a total predictive uncertainty.

These helpers are model-agnostic: they work for any model whose forward returns a tensor, a tuple of tensors, or a NamedTuple of tensors (e.g. the consistency model’s ConsistencyOutput). They are provided for downstream use and are not wired into the training/validation loops.

Functions

enable_mc_dropout(model)	Put the model in eval mode but re-enable only the dropout layers.
mc_predict(model, x[, n_samples])	Run n_samples stochastic-dropout forward passes and reduce.

Module Contents

enable_mc_dropout(model)

Put the model in eval mode but re-enable only the dropout layers.

BatchNorm / InstanceNorm and everything else stay in eval mode (using their running statistics); the dropout layers become stochastic again so repeated forward passes differ.

Parameters:

model (torch.nn.Module)

Return type:

torch.nn.Module

mc_predict(model, x, n_samples=30)

Run n_samples stochastic-dropout forward passes and reduce.

Parameters:

• model (nn.Module) – A model containing dropout layers (e.g. built with dropout > 0).

• x (input accepted by model.forward)

• n_samples (int) – Number of stochastic passes.

Returns:

The per-element mean and standard deviation across passes, with the same structure as the model output (tensor / tuple / NamedTuple). The model is left in MC-dropout mode; call model.eval() to restore.

Return type:

mean, std