ClassificationRoutine#

class torch_uncertainty.routines.ClassificationRoutine(model, num_classes, loss=None, *, is_ensemble=False, num_tta=1, format_batch_fn=None, optim_recipe=None, mixup_params=None, eval_ood=False, eval_shift=False, eval_grouping_loss=False, ood_criterion='msp', post_processing=None, num_bins_calibration_error=15, log_plots=False, save_to_csv=False, csv_filename='results.csv')[source]#

Routine for training & testing on classification tasks.

Parameters:

  • model (Module) – Model to train.

  • num_classes (int) – Number of classes.

  • loss (Module | None) – Loss function to optimize the model. Defaults to None.

  • is_ensemble (bool) – Indicates whether the model is an ensemble at test time or not. Defaults to False.

  • num_tta (int) – Number of test-time augmentations (TTA). If 1: no TTA. Defaults to 1.

  • format_batch_fn (Module | None) – Function to format the batch. Defaults to None.

  • optim_recipe (Union[Callable[[Module], Union[Optimizer, Sequence[Optimizer], tuple[Sequence[Optimizer], Sequence[Union[LRScheduler, ReduceLROnPlateau, LRSchedulerConfig]]], OptimizerConfig, OptimizerLRSchedulerConfig, Sequence[OptimizerConfig], Sequence[OptimizerLRSchedulerConfig], None]], Optimizer, Sequence[Optimizer], tuple[Sequence[Optimizer], Sequence[Union[LRScheduler, ReduceLROnPlateau, LRSchedulerConfig]]], OptimizerConfig, OptimizerLRSchedulerConfig, Sequence[OptimizerConfig], Sequence[OptimizerLRSchedulerConfig], None]) – The optimizer and optionally the scheduler to use, or a callable that returns them. Defaults to None.

  • mixup_params (dict | None) – Mixup parameters dictionary. Can include mixup type, mixup mode, distance similarity, kernel tau max, kernel tau std, mixup alpha, and cutmix alpha. If None, no mixup augmentations. Defaults to None.

  • eval_ood (bool) – Indicates whether to evaluate the OOD detection performance. Defaults to False.

  • eval_shift (bool) – Indicates whether to evaluate the Distribution shift performance. Defaults to False.

  • eval_grouping_loss (bool) – Indicates whether to evaluate the grouping loss or not. Defaults to False.

  • ood_criterion (TUOODCriterion | str) – Criterion for the binary OOD detection task. Defaults to msp, the Maximum Softmax Probability score.

  • post_processing (PostProcessing | None) – Post-processing method to train on the calibration set. No post-processing if None. Defaults to None.

  • num_bins_calibration_error (int) – Number of bins to compute calibration error metrics. Defaults to 15.

  • log_plots (bool) – Indicates whether to log plots from metrics. Defaults to False.

  • save_to_csv (bool) – Save the results in csv. Defaults to False.

  • csv_filename (str) – Name of the csv file. Defaults to "results.csv". Note that this is only used if save_to_csv is True.

Warning

You must define optim_recipe if you do not use the Lightning CLI.

Warning

When using an ensemble model, you must: 1. Set is_ensemble to True. 2. Set format_batch_fn to torch_uncertainty.transforms.RepeatTarget(num_repeats=num_estimators). 3. Ensure that the model’s forward pass outputs a tensor of shape \((M \times B, C)\), where \(M\) is the number of estimators, \(B\) is the batch size, \(C\) is the number of classes.

For automated batch handling, consider using the available model wrappers in torch_uncertainty.methods.

Note

If eval_ood is True, we perform a binary classification and update the OOD-related metrics twice: - once during the test on ID values where the given binary label is 0 (for ID) - once during the test on OOD values where the given binary label is 1 (for OOD)

Note

optim_recipe can be anything that can be returned by LightningModule.configure_optimizers(). Find more details here.

forward(inputs, save_feats=False)[source]#

Forward pass of the inner model.

Parameters:

  • inputs (Tensor) – input tensor.

  • save_feats (bool) – whether to store the features or not. Defaults to False.

Return type:

Tensor

Note

The features are stored in the self.features attribute.

on_test_epoch_end()[source]#

Compute, log, and plot the values of the collected metrics in test_step.

Return type:

None

on_test_start()[source]#

Prepare the test step.

Setup the post-processing dataset and fit the post-processing method if needed, prepares the storage lists for logit plotting and update the batchnorms if needed.

Return type:

None

on_train_start()[source]#

Put the hyperparameters in tensorboard.

Return type:

None

on_validation_epoch_end()[source]#

Compute and log the values of the collected metrics in validation_step.

Return type:

None

on_validation_start()[source]#

Prepare the validation step.

Update the model’s wrapper and the batchnorms if needed.

Return type:

None

test_step(batch, batch_idx, dataloader_idx=0)[source]#

Perform a single test step based on the input tensors.

Compute the prediction of the model and the value of the metrics on the test batch. Also handle OOD and distribution-shifted images.

Parameters:

  • batch (tuple[Tensor, Tensor]) – the test data and their corresponding targets.

  • batch_idx (int) – the number of the current batch (unused).

  • dataloader_idx (int) – 0 if in-distribution, 1 if out-of-distribution and 2 if distribution-shifted.

Return type:

None

training_step(batch)[source]#

Perform a single training step based on the input tensors.

Parameters:

batch (tuple[Tensor, Tensor]) – the training data and their corresponding targets.

Returns:

the loss corresponding to this training step.

Return type:

Tensor

validation_step(batch)[source]#

Perform a single validation step based on the input tensors.

Compute the prediction of the model and the value of the metrics on the validation batch.

Parameters:

batch (tuple[Tensor, Tensor]) – the validation data and their corresponding targets

Return type:

None