RegressionRoutine

class torch_uncertainty.routines.RegressionRoutine(model, output_dim, probabilistic, loss, num_estimators=1, optim_recipe=None, format_batch_fn=None)

Routine for efficient training and testing on regression tasks using LightningModule.

Parameters:

• model (torch.nn.Module) – Model to train.

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

• probabilistic (bool) – Whether the model is probabilistic, i.e., outputs a PyTorch distribution.

• loss (torch.nn.Module) – Loss function to optimize the model.

• num_estimators (int, optional) – The number of estimators for the ensemble. Defaults to 1 (single model).

• optim_recipe (dict or torch.optim.Optimizer, optional) – The optimizer and optionally the scheduler to use. Defaults to None.

• format_batch_fn (torch.nn.Module, optional) – The function to format the batch. Defaults to None.

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.

forward(inputs)

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.

Return type:

Tensor