scvi.model.base.PyroJitGuideWarmup#

class scvi.model.base.PyroJitGuideWarmup(dataloader=None)[source]#

A callback to warmup a Pyro guide.

This helps initialize all the relevant parameters by running one minibatch through the Pyro model.

Attributes table#

state_key

Identifier for the state of the callback.

Methods table#

on_after_backward(trainer, pl_module)

Called after loss.backward() and before optimizers are stepped.

on_batch_end(trainer, pl_module)

Called when the training batch ends.

on_batch_start(trainer, pl_module)

Called when the training batch begins.

on_before_accelerator_backend_setup(trainer, ...)

Called before accelerator is being setup.

on_before_backward(trainer, pl_module, loss)

Called before loss.backward().

on_before_optimizer_step(trainer, pl_module, ...)

Called before optimizer.step().

on_before_zero_grad(trainer, pl_module, ...)

Called before optimizer.zero_grad().

on_configure_sharded_model(trainer, pl_module)

Called before configure sharded model.

on_epoch_end(trainer, pl_module)

Called when either of train/val/test epoch ends.

on_epoch_start(trainer, pl_module)

Called when either of train/val/test epoch begins.

on_exception(trainer, pl_module, exception)

Called when any trainer execution is interrupted by an exception.

on_fit_end(trainer, pl_module)

Called when fit ends.

on_fit_start(trainer, pl_module)

Called when fit begins.

on_init_end(trainer)

Called when the trainer initialization ends, model has not yet been set.

on_init_start(trainer)

Called when the trainer initialization begins, model has not yet been set.

on_keyboard_interrupt(trainer, pl_module)

Deprecated since version v1.5.

on_load_checkpoint(trainer, pl_module, ...)

Called when loading a model checkpoint, use to reload state.

on_predict_batch_end(trainer, pl_module, ...)

Called when the predict batch ends.

on_predict_batch_start(trainer, pl_module, ...)

Called when the predict batch begins.

on_predict_end(trainer, pl_module)

Called when predict ends.

on_predict_epoch_end(trainer, pl_module, outputs)

Called when the predict epoch ends.

on_predict_epoch_start(trainer, pl_module)

Called when the predict epoch begins.

on_predict_start(trainer, pl_module)

Called when the predict begins.

on_pretrain_routine_end(trainer, pl_module)

Called when the pretrain routine ends.

on_pretrain_routine_start(trainer, pl_module)

Called when the pretrain routine begins.

on_sanity_check_end(trainer, pl_module)

Called when the validation sanity check ends.

on_sanity_check_start(trainer, pl_module)

Called when the validation sanity check starts.

on_save_checkpoint(trainer, pl_module, ...)

Called when saving a model checkpoint, use to persist state.

on_test_batch_end(trainer, pl_module, ...)

Called when the test batch ends.

on_test_batch_start(trainer, pl_module, ...)

Called when the test batch begins.

on_test_end(trainer, pl_module)

Called when the test ends.

on_test_epoch_end(trainer, pl_module)

Called when the test epoch ends.

on_test_epoch_start(trainer, pl_module)

Called when the test epoch begins.

on_test_start(trainer, pl_module)

Called when the test begins.

on_train_batch_end(trainer, pl_module, ...)

Called when the train batch ends.

on_train_batch_start(trainer, pl_module, ...)

Called when the train batch begins.

on_train_end(trainer, pl_module)

Called when the train ends.

on_train_epoch_end(trainer, pl_module)

Called when the train epoch ends.

on_train_epoch_start(trainer, pl_module)

Called when the train epoch begins.

on_train_start(trainer, pl_module)

Way to warmup Pyro Guide in an automated way.

on_validation_batch_end(trainer, pl_module, ...)

Called when the validation batch ends.

on_validation_batch_start(trainer, ...)

Called when the validation batch begins.

on_validation_end(trainer, pl_module)

Called when the validation loop ends.

on_validation_epoch_end(trainer, pl_module)

Called when the val epoch ends.

on_validation_epoch_start(trainer, pl_module)

Called when the val epoch begins.

on_validation_start(trainer, pl_module)

Called when the validation loop begins.

setup(trainer, pl_module[, stage])

Called when fit, validate, test, predict, or tune begins.

teardown(trainer, pl_module[, stage])

Called when fit, validate, test, predict, or tune ends.

Attributes#

state_key#

PyroJitGuideWarmup.state_key#

Identifier for the state of the callback.

Used to store and retrieve a callback’s state from the checkpoint dictionary by checkpoint["callbacks"][state_key]. Implementations of a callback need to provide a unique state key if 1) the callback has state and 2) it is desired to maintain the state of multiple instances of that callback.

Return type

str

Methods#

on_after_backward#

PyroJitGuideWarmup.on_after_backward(trainer, pl_module)#

Called after loss.backward() and before optimizers are stepped.

Return type

None

on_batch_end#

PyroJitGuideWarmup.on_batch_end(trainer, pl_module)#

Called when the training batch ends.

Return type

None

on_batch_start#

PyroJitGuideWarmup.on_batch_start(trainer, pl_module)#

Called when the training batch begins.

Return type

None

on_before_accelerator_backend_setup#

PyroJitGuideWarmup.on_before_accelerator_backend_setup(trainer, pl_module)#

Called before accelerator is being setup.

Return type

None

on_before_backward#

PyroJitGuideWarmup.on_before_backward(trainer, pl_module, loss)#

Called before loss.backward().

Return type

None

on_before_optimizer_step#

PyroJitGuideWarmup.on_before_optimizer_step(trainer, pl_module, optimizer, opt_idx)#

Called before optimizer.step().

Return type

None

on_before_zero_grad#

PyroJitGuideWarmup.on_before_zero_grad(trainer, pl_module, optimizer)#

Called before optimizer.zero_grad().

Return type

None

on_configure_sharded_model#

PyroJitGuideWarmup.on_configure_sharded_model(trainer, pl_module)#

Called before configure sharded model.

Return type

None

on_epoch_end#

PyroJitGuideWarmup.on_epoch_end(trainer, pl_module)#

Called when either of train/val/test epoch ends.

Return type

None

on_epoch_start#

PyroJitGuideWarmup.on_epoch_start(trainer, pl_module)#

Called when either of train/val/test epoch begins.

Return type

None

on_exception#

PyroJitGuideWarmup.on_exception(trainer, pl_module, exception)#

Called when any trainer execution is interrupted by an exception.

Return type

None

on_fit_end#

PyroJitGuideWarmup.on_fit_end(trainer, pl_module)#

Called when fit ends.

Return type

None

on_fit_start#

PyroJitGuideWarmup.on_fit_start(trainer, pl_module)#

Called when fit begins.

Return type

None

on_init_end#

PyroJitGuideWarmup.on_init_end(trainer)#

Called when the trainer initialization ends, model has not yet been set.

Return type

None

on_init_start#

PyroJitGuideWarmup.on_init_start(trainer)#

Called when the trainer initialization begins, model has not yet been set.

Return type

None

on_keyboard_interrupt#

PyroJitGuideWarmup.on_keyboard_interrupt(trainer, pl_module)#

Deprecated since version v1.5: This callback hook was deprecated in v1.5 in favor of on_exception and will be removed in v1.7.

Called when any trainer execution is interrupted by KeyboardInterrupt.

Return type

None

on_load_checkpoint#

PyroJitGuideWarmup.on_load_checkpoint(trainer, pl_module, callback_state)#

Called when loading a model checkpoint, use to reload state.

Parameters
trainer : Trainer

the current Trainer instance.

pl_module : LightningModule

the current LightningModule instance.

callback_state : {str: Any}Dict[str, Any]

the callback state returned by on_save_checkpoint.

Note

The on_load_checkpoint won’t be called with an undefined state. If your on_load_checkpoint hook behavior doesn’t rely on a state, you will still need to override on_save_checkpoint to return a dummy state.

Return type

None

on_predict_batch_end#

PyroJitGuideWarmup.on_predict_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)#

Called when the predict batch ends.

Return type

None

on_predict_batch_start#

PyroJitGuideWarmup.on_predict_batch_start(trainer, pl_module, batch, batch_idx, dataloader_idx)#

Called when the predict batch begins.

Return type

None

on_predict_end#

PyroJitGuideWarmup.on_predict_end(trainer, pl_module)#

Called when predict ends.

Return type

None

on_predict_epoch_end#

PyroJitGuideWarmup.on_predict_epoch_end(trainer, pl_module, outputs)#

Called when the predict epoch ends.

Return type

None

on_predict_epoch_start#

PyroJitGuideWarmup.on_predict_epoch_start(trainer, pl_module)#

Called when the predict epoch begins.

Return type

None

on_predict_start#

PyroJitGuideWarmup.on_predict_start(trainer, pl_module)#

Called when the predict begins.

Return type

None

on_pretrain_routine_end#

PyroJitGuideWarmup.on_pretrain_routine_end(trainer, pl_module)#

Called when the pretrain routine ends.

Return type

None

on_pretrain_routine_start#

PyroJitGuideWarmup.on_pretrain_routine_start(trainer, pl_module)#

Called when the pretrain routine begins.

Return type

None

on_sanity_check_end#

PyroJitGuideWarmup.on_sanity_check_end(trainer, pl_module)#

Called when the validation sanity check ends.

Return type

None

on_sanity_check_start#

PyroJitGuideWarmup.on_sanity_check_start(trainer, pl_module)#

Called when the validation sanity check starts.

Return type

None

on_save_checkpoint#

PyroJitGuideWarmup.on_save_checkpoint(trainer, pl_module, checkpoint)#

Called when saving a model checkpoint, use to persist state.

Parameters
trainer : Trainer

the current Trainer instance.

pl_module : LightningModule

the current LightningModule instance.

checkpoint : {str: Any}Dict[str, Any]

the checkpoint dictionary that will be saved.

Return type

dict

Returns

The callback state.

on_test_batch_end#

PyroJitGuideWarmup.on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)#

Called when the test batch ends.

Return type

None

on_test_batch_start#

PyroJitGuideWarmup.on_test_batch_start(trainer, pl_module, batch, batch_idx, dataloader_idx)#

Called when the test batch begins.

Return type

None

on_test_end#

PyroJitGuideWarmup.on_test_end(trainer, pl_module)#

Called when the test ends.

Return type

None

on_test_epoch_end#

PyroJitGuideWarmup.on_test_epoch_end(trainer, pl_module)#

Called when the test epoch ends.

Return type

None

on_test_epoch_start#

PyroJitGuideWarmup.on_test_epoch_start(trainer, pl_module)#

Called when the test epoch begins.

Return type

None

on_test_start#

PyroJitGuideWarmup.on_test_start(trainer, pl_module)#

Called when the test begins.

Return type

None

on_train_batch_end#

PyroJitGuideWarmup.on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, unused=0)#

Called when the train batch ends.

Return type

None

on_train_batch_start#

PyroJitGuideWarmup.on_train_batch_start(trainer, pl_module, batch, batch_idx, unused=0)#

Called when the train batch begins.

Return type

None

on_train_end#

PyroJitGuideWarmup.on_train_end(trainer, pl_module)#

Called when the train ends.

Return type

None

on_train_epoch_end#

PyroJitGuideWarmup.on_train_epoch_end(trainer, pl_module)#

Called when the train epoch ends.

To access all batch outputs at the end of the epoch, either:

  1. Implement training_epoch_end in the LightningModule and access outputs via the module OR

  2. Cache data across train batch hooks inside the callback implementation to post-process in this hook.

Return type

None

on_train_epoch_start#

PyroJitGuideWarmup.on_train_epoch_start(trainer, pl_module)#

Called when the train epoch begins.

Return type

None

on_train_start#

PyroJitGuideWarmup.on_train_start(trainer, pl_module)[source]#

Way to warmup Pyro Guide in an automated way.

Also device agnostic.

on_validation_batch_end#

PyroJitGuideWarmup.on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)#

Called when the validation batch ends.

Return type

None

on_validation_batch_start#

PyroJitGuideWarmup.on_validation_batch_start(trainer, pl_module, batch, batch_idx, dataloader_idx)#

Called when the validation batch begins.

Return type

None

on_validation_end#

PyroJitGuideWarmup.on_validation_end(trainer, pl_module)#

Called when the validation loop ends.

Return type

None

on_validation_epoch_end#

PyroJitGuideWarmup.on_validation_epoch_end(trainer, pl_module)#

Called when the val epoch ends.

Return type

None

on_validation_epoch_start#

PyroJitGuideWarmup.on_validation_epoch_start(trainer, pl_module)#

Called when the val epoch begins.

Return type

None

on_validation_start#

PyroJitGuideWarmup.on_validation_start(trainer, pl_module)#

Called when the validation loop begins.

Return type

None

setup#

PyroJitGuideWarmup.setup(trainer, pl_module, stage=None)#

Called when fit, validate, test, predict, or tune begins.

Return type

None

teardown#

PyroJitGuideWarmup.teardown(trainer, pl_module, stage=None)#

Called when fit, validate, test, predict, or tune ends.

Return type

None