sage.utils.checkpoint

Filename : checkpoint.py Description : Short description of the file

Created on 2026-03-22 11:10:40

__author__ = Narenraju Nagarajan __copyright__ = Copyright 2026, ProjectName __license__ = MIT Licence __version__ = 0.0.1 __maintainer__ = Narenraju Nagarajan __affiliation__ = N/A __email__ = N/A __status__ = [‘inProgress’, ‘Archived’, ‘inUsage’, ‘Debugging’]

GitHub Repository: NULL

Documentation: NULL

Classes

CheckpointManager

Manages saving and loading of training checkpoints.

Module Contents

class CheckpointManager(cfg, data_cfg, model, optimizer, scheduler, scaler)

Manages saving and loading of training checkpoints.

On construction, creates the checkpoint directory and writes one-time JSON snapshots of cfg and data_cfg so that the hyperparameters that produced each checkpoint are always recoverable.

Three checkpoint types are maintained:

• latest.pt — overwritten every call to save(); safe resume point after a crash or pre-emption.

• epoch_{N}.pt — per-epoch copy (optional; controlled by save_epoch_ckpt). Allows rolling back to any specific epoch.

• best.pt — copy of latest.pt whenever val_loss improves; the recommended checkpoint for inference.

Each checkpoint file contains the full training state: model weights, optimiser and scheduler states, AMP scaler state, configurations, and all four RNG states (Python, NumPy, CPU torch, CUDA torch) for bit-exact reproducibility.

Parameters:

• cfg (BaseConfig) – Training configuration (provides export_dir, autocast, dtype).

• data_cfg (BaseDataConfig) – Data configuration (saved to JSON snapshot only).

• model (nn.Module) – The model being trained (may be a torch.compile wrapper).

• optimizer (torch.optim.Optimizer)

• scheduler (torch.optim.lr_scheduler._LRScheduler)

• scaler (torch.amp.GradScaler)

cfg

data_cfg

model

optimizer

scheduler

scaler

best_val_loss

ckpt_dir

latest_path

best_path

save(epoch, val_loss=None, save_epoch_ckpt=True)

Save the current training state to disk.

Always writes latest.pt. If val_loss is better than the current best, also copies to best.pt.

Parameters:

• epoch (int) – Current epoch (0-based).

• val_loss (float or None) – Validation loss; used to decide whether to update best.pt.

• save_epoch_ckpt (bool) – If True (default), also write epoch_{epoch}.pt.

load_latest(map_location='cpu')

Load latest snapshot and resume training

Parameters:

map_location (str, optional) – _description_. Defaults to “cpu”.

Returns:

_description_

Return type:

_type_

Usage:

start_epoch = ckpt_mgr.load_latest(map_location=cfg.device) for nepoch in range(start_epoch, cfg.num_epochs): …

load_best(map_location='cpu')

Load best snapshot for inference

Parameters:

map_location (str, optional) – _description_. Defaults to “cpu”.

Returns:

_description_

Return type:

_type_

Usage:

ckpt_mgr.load_best(map_location=cfg.device) model.eval()

load_epoch(epoch, map_location='cpu')

Load a specific per-epoch checkpoint.

Parameters:

• epoch (int) – Epoch index of the checkpoint to load.

• map_location (str) – Device mapping for torch.load.

Returns:

epoch + 1 — the next epoch to run.

Return type:

int