RegressionRoutine#

class torch_uncertainty.routines.RegressionRoutine(model, output_dim, loss=None, dist_family=None, dist_estimate='mean', *, is_ensemble=False, optim_recipe=None, eval_shift=False, format_batch_fn=None, log_plots=False, num_bins_calibration_error=15, save_to_csv=False, csv_filename='results.csv')[source]#

Routine for training & testing on regression tasks.

Parameters:

  • model (Module) – Model to train.

  • output_dim (int) – Number of outputs of the model.

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

  • dist_family (str | None) – Distribution family to use for probabilistic regression. If None, performs point-wise regression. Defaults to None.

  • dist_estimate (str | DistEstimate) – The estimate to use when computing point-wise metrics. Defaults to mean.

  • is_ensemble (bool) – Whether the model is an ensemble. Defaults to False.

  • 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.

  • eval_shift (bool) – Whether to evaluate distribution-shift performance. Defaults to False.

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

  • log_plots (bool) – Whether to log figures in the logger. Defaults to False.

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

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

  • csv_filename (str) – Name of the CSV file. Defaults to results.csv. Used only when save_to_csv is True.

Warning

If probabilistic is True, the model must output a PyTorch distribution.

Warning

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

Note

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

evaluation_forward(inputs)[source]#

Get the prediction and handle predicted eventual distribution parameters.

Parameters:

inputs (Tensor) – The input data.

Returns:

the prediction as a Tensor and a distribution.

Return type:

tuple[Tensor, Distribution | None]

forward(inputs)[source]#

Forward pass of the routine.

The forward pass automatically squeezes the output if the regression is one-dimensional and if the routine contains a single model.

Parameters:

inputs (Tensor) – The input tensor.

Returns:

The output tensor or the parameters of the output

distribution.

Return type:

Tensor | dict[str, Tensor]

on_test_epoch_end()[source]#

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

Return type:

None

on_test_start()[source]#

Prepare the test step.

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]) – Tuple of test inputs and targets.

  • batch_idx (int) – Index of the batch in the dataloader (unused here).

  • dataloader_idx (int) – 0 for in-distribution, 1 for out-of-distribution.

Return type:

None

training_step(batch)[source]#

Perform a single training step based on the input tensors.

Parameters:

batch (tuple[Tensor, Tensor]) – Tuple of training inputs and 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]) – Tuple of validation inputs and targets.

Return type:

None