scvi.train.PyroTrainingPlan

class scvi.train.PyroTrainingPlan(pyro_module, loss_fn=None, optim=None, optim_kwargs=None, n_steps_kl_warmup=None, n_epochs_kl_warmup=400, scale_elbo=1.0, blocked=None, max_kl_weight=1.0, min_kl_weight=1e-06)

Bases: LowLevelPyroTrainingPlan

Lightning module task to train Pyro scvi-tools modules.

Parameters:

• pyro_module (PyroBaseModuleClass) – An instance of PyroBaseModuleClass. This object should have callable model and guide attributes or methods.

• loss_fn (ELBO | None (default: None)) – A Pyro loss. Should be a subclass of ELBO. If None, defaults to Trace_ELBO.

• optim (PyroOptim | None (default: None)) – A Pyro optimizer instance, e.g., Adam(). If None, defaults to Adam() optimizer with a learning rate of 1e-3.

• optim_kwargs (dict | None (default: None)) – Keyword arguments for default optimiser Adam().

• n_steps_kl_warmup (int | None (default: None)) – Number of training steps (minibatches) to scale weight on KL divergences from 0 to 1. Only activated when n_epochs_kl_warmup is set to None.

• n_epochs_kl_warmup (int | None (default: 400)) – Number of epochs to scale weight on KL divergences from 0 to 1. Overrides n_steps_kl_warmup when both are not None.

• scale_elbo (float (default: 1.0)) – Scale ELBO using scale(). Potentially useful for avoiding numerical inaccuracy when working with very large ELBO.

• blocked (list | None (default: None)) – A list of Pyro parameters to block during training. If None, defaults to train all parameters.

• min_kl_weight (float (default: 1e-06)) – Minimum KL weight during warmup. Defaults to 1e-6.

• max_kl_weight (float (default: 1.0)) – Maximum KL weight during warmup. Defaults to 1.0.

Attributes table

training

Methods table

backward(*args, **kwargs)	No-op, as Pyro computes the gradients inside the training step.
configure_optimizers()	Shim optimizer for PyTorch Lightning.
optimizer_step(*args, **kwargs)	No-op, as the Pyro optimizer steps inside the training step.
training_step(batch, batch_idx)	Training step for Pyro training.

Attributes

PyroTrainingPlan.training: bool

Methods

PyroTrainingPlan.backward(*args, **kwargs)

No-op, as Pyro computes the gradients inside the training step.

PyroTrainingPlan.configure_optimizers()

Shim optimizer for PyTorch Lightning.

PyTorch Lightning wants to take steps on an optimizer returned by this function in order to increment the global step count. See PyTorch Lightning optimizer manual loop.

Here we provide a shim optimizer that we can take steps on at minimal computational cost in order to keep Lightning happy :).

PyroTrainingPlan.optimizer_step(*args, **kwargs)

No-op, as the Pyro optimizer steps inside the training step.

PyroTrainingPlan.training_step(batch, batch_idx)

Training step for Pyro training.