deepextractor.utils.mc_dropout

Monte Carlo Dropout utilities for uncertainty estimation.

At test time, keeping Dropout active and running N stochastic forward passes gives N samples from an approximate posterior over model outputs. Works with both UNET2D (Dropout2d, STFT domain) and UNET1D (Dropout1d, time domain).

For STFT-domain models (UNET2D), work in time domain, not STFT domain. The model outputs magnitude+phase spectrograms. Phase is a circular quantity, so naive averaging across passes in the STFT domain is incorrect (e.g. phases of +170 deg and -170 deg average to 0 deg, not +-180 deg). The right approach is to apply iSTFT to each sample before aggregating, so the posterior is over time-domain reconstructions — which is also the output space you care about physically. UNET1D has no such concern since it already outputs time-domain samples directly.

Usage (STFT / UNET2D):

from deepextractor.utils import enable_mc_dropout, mc_predict
from deepextractor.utils.stft import apply_istft

model.eval()              # freeze BatchNorm running stats
enable_mc_dropout(model)  # re-enable Dropout2d only

import functools
istft_fn = functools.partial(
    apply_istft, n_fft=512, hop_length=32, win_length=64,
    window=torch.hann_window(64),
)

# Full posterior in time domain — shape (n_passes, N, L)
samples = mc_predict(model, x, n_passes=100, postprocess_fn=istft_fn)

mean   = samples.mean(dim=0)            # shape (N, L)
median = samples.median(dim=0).values
std    = samples.std(dim=0)             # per-sample uncertainty vs time
lo, hi = samples.quantile(torch.tensor([0.05, 0.95]), dim=0)

# Or get a summary directly
mean = mc_predict(model, x, n_passes=100, postprocess_fn=istft_fn, reduction="mean")

Usage (time domain / UNET1D) — no postprocess_fn needed, model output is already time domain:

model.eval()
enable_mc_dropout(model)
samples = mc_predict(model, x, n_passes=100)   # (n_passes, N, C, L)

Module Contents

deepextractor.utils.mc_dropout.enable_mc_dropout(model: torch.nn.Module) → None[source]

Set all Dropout1d/Dropout2d modules to train mode while leaving everything else eval.

Call this after model.eval() to activate Monte Carlo Dropout at inference time. BatchNorm layers stay in eval mode so they use their running statistics rather than per-batch statistics. Works for both UNET2D (Dropout2d) and UNET1D (Dropout1d).

deepextractor.utils.mc_dropout.mc_predict(model: torch.nn.Module, x: torch.Tensor, n_passes: int = 100, postprocess_fn: Callable | None = None, reduction: str = 'none') → torch.Tensor[source]

Draw n_passes stochastic samples from the approximate posterior.

Each forward pass uses a different dropout mask. An optional postprocess_fn is applied to each raw model output before stacking — use this to convert from STFT domain to time domain via iSTFT (recommended for UNET2D, see module docstring for why averaging in STFT domain is incorrect). Not needed for UNET1D, whose output is already time domain.

Parameters:

  • model (nn.Module with Dropout1d/Dropout2d layers already in train) – mode (call enable_mc_dropout first).

  • x (torch.Tensor — input batch, any rank (N, C, L) for UNET1D) – or (N, C, H, W) for UNET2D; the batch dimension (dim 0) is what gets repeated across passes.

  • n_passes (int — number of stochastic forward passes (50-100 typical).)

  • postprocess_fn (optional callable applied to each model output before) – stacking, e.g. functools.partial(apply_istft, ...) to convert STFT outputs to time-domain signals.

  • reduction (str — how to summarise the stacked samples:) –

    "none" — return all samples (recommended for

    forming a full posterior)

    "mean" — sample-wise mean "median" — sample-wise median

Returns:

If reduction="none": shape (n_passes, N, ...) where ... is the shape of a single (optionally post-processed) output. If reduction="mean" or "median": shape (N, ...).

Return type:

torch.Tensor

Examples

>>> samples = mc_predict(model, x, n_passes=100, postprocess_fn=istft_fn)
>>> mean   = samples.mean(dim=0)
>>> median = samples.median(dim=0).values
>>> std    = samples.std(dim=0)    # uncertainty map over time
>>> lo, hi = samples.quantile(torch.tensor([0.05, 0.95]), dim=0)