def test_save_best_state_callback(save_path):
n_latent = 5
adata = synthetic_iid()
SCVI.setup_anndata(adata, batch_key="batch", labels_key="labels")
model = SCVI(adata, n_latent=n_latent)
callbacks = [SaveBestState(verbose=True)]
model.train(3, check_val_every_n_epoch=1, train_size=0.5, callbacks=callbacks)
def test_save_best_state_callback(save_path):
n_latent = 5
adata = synthetic_iid()
model = SCVI(adata, n_latent=n_latent)
callbacks = [SaveBestState(verbose=True)]
model.train(3,
check_val_every_n_epoch=1,
train_size=0.5,
callbacks=callbacks)
def train(
self,
max_epochs: int = 500,
lr: float = 1e-4,
use_gpu: Optional[Union[str, int, bool]] = None,
train_size: float = 0.9,
validation_size: Optional[float] = None,
batch_size: int = 128,
weight_decay: float = 1e-3,
eps: float = 1e-08,
early_stopping: bool = True,
early_stopping_patience: int = 50,
save_best: bool = True,
check_val_every_n_epoch: Optional[int] = None,
n_steps_kl_warmup: Union[int, None] = None,
n_epochs_kl_warmup: Union[int, None] = 50,
plan_kwargs: Optional[dict] = None,
**kwargs,
):
"""
Trains the model using amortized variational inference.
Parameters
----------
max_epochs
Number of passes through the dataset.
lr
Learning rate for optimization.
use_gpu
Use default GPU if available (if None or True), or index of GPU to use (if int),
or name of GPU (if str, e.g., `'cuda:0'`), or use CPU (if False).
train_size
Size of training set in the range [0.0, 1.0].
validation_size
Size of the test set. If `None`, defaults to 1 - `train_size`. If
`train_size + validation_size < 1`, the remaining cells belong to a test set.
batch_size
Minibatch size to use during training.
weight_decay
weight decay regularization term for optimization
eps
Optimizer eps
early_stopping
Whether to perform early stopping with respect to the validation set.
early_stopping_patience
How many epochs to wait for improvement before early stopping
save_best
Save the best model state with respect to the validation loss (default), or use the final
state in the training procedure
check_val_every_n_epoch
Check val every n train epochs. By default, val is not checked, unless `early_stopping` is `True`.
If so, val is checked every epoch.
n_steps_kl_warmup
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. If `None`, defaults
to `floor(0.75 * adata.n_obs)`.
n_epochs_kl_warmup
Number of epochs to scale weight on KL divergences from 0 to 1.
Overrides `n_steps_kl_warmup` when both are not `None`.
plan_kwargs
Keyword args for :class:`~scvi.train.TrainingPlan`. Keyword arguments passed to
`train()` will overwrite values present in `plan_kwargs`, when appropriate.
**kwargs
Other keyword args for :class:`~scvi.train.Trainer`.
"""
update_dict = dict(
lr=lr,
weight_decay=weight_decay,
eps=eps,
n_epochs_kl_warmup=n_epochs_kl_warmup,
n_steps_kl_warmup=n_steps_kl_warmup,
optimizer="AdamW",
)
if plan_kwargs is not None:
plan_kwargs.update(update_dict)
else:
plan_kwargs = update_dict
if save_best:
if "callbacks" not in kwargs.keys():
kwargs["callbacks"] = []
kwargs["callbacks"].append(
SaveBestState(monitor="reconstruction_loss_validation"))
super().train(
max_epochs=max_epochs,
train_size=train_size,
use_gpu=use_gpu,
validation_size=validation_size,
early_stopping=early_stopping,
early_stopping_monitor="reconstruction_loss_validation",
early_stopping_patience=early_stopping_patience,
plan_kwargs=plan_kwargs,
check_val_every_n_epoch=check_val_every_n_epoch,
batch_size=batch_size,
**kwargs,
)