RegressionRoutine¶

class torch_uncertainty.routines.RegressionRoutine(model, output_dim, probabilistic, loss, is_ensemble=False, optim_recipe=None, eval_shift=False, format_batch_fn=None)[source]¶

Routine for training & testing on regression tasks.

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.
is_ensemble (bool, optional) – Whether the model is an ensemble. Defaults to False.
optim_recipe (dict or torch.optim.Optimizer, optional) – The optimizer and optionally the scheduler to use. Defaults to None.
eval_shift (bool, optional) – Indicates whether to evaluate the Distribution shift performance. Defaults to False.
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)[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.
Return type:: Tensor