Shortcuts

PixelRegressionRoutine

class torch_uncertainty.routines.PixelRegressionRoutine(model, output_dim, probabilistic, loss, is_ensemble=False, format_batch_fn=None, optim_recipe=None, eval_shift=False, num_image_plot=4, log_plots=False)[source]

Routine for training & testing on pixel regression tasks.

Parameters:
  • model (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 (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 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 (nn.Module, optional) – The function to format the batch. Defaults to None.

  • num_image_plot (int, optional) – Number of images to plot. Defaults to 4.

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

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