BayesianModelWrapper¶

class BayesianModelWrapper(base_model, n_samples=30, class_dim=-1, *, inject_dropout=True, dropout_p=0.2)[source]¶

Wrap a TIAToolbox model and add Monte Carlo Dropout at infer_batch.

The wrapped model must contain at least one Dropout layer. For pretrained models that ship without Dropout (e.g. resnet18-kather100k, alexnet-kather100k), pass inject_dropout=True (default) and a nn.Dropout is inserted before the linear classifier automatically, via inject_dropout_before_classifier.

After each infer_batch call, the per-batch uncertainty decomposition (mean_probs, aleatoric, epistemic, total, mc_probs) is appended to self.uncertainty_stats in inference order, so it can be aligned with the coordinates array written to the engine’s output .zarr.

Example

Wrap a pretrained patch classifier and run inference with the patch predictor, collecting uncertainty per patch:

>>> from tiatoolbox.models.architecture import get_pretrained_model
>>> from tiatoolbox.models.engine.patch_predictor import PatchPredictor
>>> from tiatoolbox.models.architecture.bayesian_wrapper import (
...     BayesianModelWrapper,
... )
>>>
>>> model, _ = get_pretrained_model("resnet18-kather100k")
>>> bayesian_model = BayesianModelWrapper(
...     base_model=model,
...     n_samples=30,  # number of MC forward passes
...     inject_dropout=True,  # model ships without Dropout
...     dropout_p=0.2,
... )
>>>
>>> predictor = PatchPredictor(model=bayesian_model)
>>> outputs = predictor.run(images=["sample_img.png"], patch_mode=True)
>>>
>>> # After predict() finishes, concatenate the per-batch stats:
>>> epistemic = np.concatenate(
...     [s["epistemic"] for s in bayesian_model.uncertainty_stats], axis=0
... )  # one value per patch, aligned with outputs["coordinates"]
>>> total = np.concatenate(
...     [s["total"] for s in bayesian_model.uncertainty_stats], axis=0
... )

Note

The wrapped model must be a classification model whose infer_batch(model, images, device) returns per-sample class probabilities of shape [B, C] (e.g. CNNModel, TimmModel).

Initialize BayesianModelWrapper.

Parameters:

  • base_model (torch.nn.Module) – TIAToolbox model to wrap. Must expose an infer_batch static method matching (model, images, device).

  • n_samples (int) – Number of Monte Carlo forward passes per batch.

  • class_dim (int) – Class axis in the per-sample probabilities (excluding T).

  • inject_dropout (bool) – If True and the model has no Dropout layers, insert one before the linear classifier automatically.

  • dropout_p (float) – Dropout probability used by inject_dropout.

Raises:

  • ValueError – If the model has no Dropout layers and inject_dropout is False.

  • RuntimeError – If dropout injection was requested but failed.

Methods

infer_batch

Run n_samples MC forward passes under mc_dropout_mode.

Attributes

training

static infer_batch(model, batch_data, device='cpu')[source]¶

Run n_samples MC forward passes under mc_dropout_mode.

Returns the MC-mean per-class probabilities of shape [B, C] (numpy), and appends a per-batch uncertainty dict to model.uncertainty_stats.

Parameters:
Return type:

np.ndarray