urbanlc.model.train_utils.save_checkpoint

urbanlc.model.train_utils.save_checkpoint(save_dir: str, name: str, model: Module, optimizer: Optimizer, scheduler: _LRScheduler, epoch: int) → str

Save model checkpoint, including model state, optimizer state, scheduler state, and epoch.

This function saves a model checkpoint to a file in the specified directory. The checkpoint includes the model state, optimizer state, scheduler state, and current epoch.

Parameters:

• save_dir (str) – Directory to save the checkpoint file.

• name (str) – Name of the checkpoint file.

• model (nn.Module) – PyTorch model to be saved.

• optimizer (torch.optim.Optimizer) – Optimizer used during training.

• scheduler (torch.optim.lr_scheduler._LRScheduler) – Learning rate scheduler. Defaults to None.

• epoch (int) – Current epoch.

Returns:

Filepath of the saved checkpoint.

Return type:

str