Skip to content

rslearn.train.callbacks.checkpointing

checkpointing

Checkpoint callbacks for rslearn.

BestLastCheckpoint

Bases: Callback

Saves both last.ckpt (every validation) and best.ckpt (when metric improves).

Source code in rslearn/train/callbacks/checkpointing.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
class BestLastCheckpoint(Callback):
    """Saves both last.ckpt (every validation) and best.ckpt (when metric improves)."""

    def __init__(
        self,
        dirpath: str,
        monitor: str = "val_loss",
        mode: Literal["min", "max"] = "min",
    ) -> None:
        """Create a new BestLastCheckpoint.

        Args:
            dirpath: directory to save checkpoints in (local or cloud path).
            monitor: metric key to monitor for best checkpoint.
            mode: "min" if lower is better, "max" if higher is better.
        """
        self.dirpath = dirpath
        self.monitor = monitor
        self.mode = mode
        self._best_value: float | None = None

    def state_dict(self) -> dict[str, Any]:
        """Persist best value so it survives checkpoint resume."""
        return {"best_value": self._best_value}

    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        """Restore best value from checkpoint on resume."""
        self._best_value = state_dict.get("best_value")

    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
        """Save last.ckpt always, and best.ckpt when the monitored metric improves."""
        if trainer.sanity_checking:
            return

        # Check metric value to handle best.ckpt.
        if self.monitor not in trainer.callback_metrics:
            raise ValueError(
                f"did not find {self.monitor} metric (keys are {trainer.callback_metrics.keys()})"
            )
        current_val = trainer.callback_metrics[self.monitor].item()

        is_better = self._best_value is None or (
            (self.mode == "min" and current_val < self._best_value)
            or (self.mode == "max" and current_val > self._best_value)
        )
        if is_better:
            self._best_value = current_val
            best_path = os.path.join(self.dirpath, "best.ckpt")
            trainer.save_checkpoint(best_path)

            # Previously this entire function was decorated with @rank_zero_only, but this resulted
            # in NCCL timeout because all ranks must call trainer.save_checkpoint. Log only on rank 0.
            if trainer.is_global_zero:
                logger.info(
                    "saved best checkpoint to %s (%s=%s)",
                    best_path,
                    self.monitor,
                    current_val,
                )

        # Save last.ckpt after best.ckpt so that auto-resuming from last.ckpt is always
        # correct.
        # - If we crash after writing best.ckpt, we will auto-resume from previous epoch
        #   and write the same best.ckpt again.
        # - If we crash after writing last.ckpt, we will restore the correct
        #   self._best_value and won't override the checkpoint incorrectly.
        last_path = os.path.join(self.dirpath, "last.ckpt")
        trainer.save_checkpoint(last_path)
        if trainer.is_global_zero:
            logger.info("saved checkpoint to %s", last_path)

state_dict 

state_dict() -> dict[str, Any]

Persist best value so it survives checkpoint resume.

Source code in rslearn/train/callbacks/checkpointing.py 
35
36
37
def state_dict(self) -> dict[str, Any]:
    """Persist best value so it survives checkpoint resume."""
    return {"best_value": self._best_value}

load_state_dict 

load_state_dict(state_dict: dict[str, Any]) -> None

Restore best value from checkpoint on resume.

Source code in rslearn/train/callbacks/checkpointing.py 
39
40
41
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Restore best value from checkpoint on resume."""
    self._best_value = state_dict.get("best_value")

on_validation_end 

on_validation_end(trainer: Trainer, pl_module: LightningModule) -> None

Save last.ckpt always, and best.ckpt when the monitored metric improves.

Source code in rslearn/train/callbacks/checkpointing.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
    """Save last.ckpt always, and best.ckpt when the monitored metric improves."""
    if trainer.sanity_checking:
        return

    # Check metric value to handle best.ckpt.
    if self.monitor not in trainer.callback_metrics:
        raise ValueError(
            f"did not find {self.monitor} metric (keys are {trainer.callback_metrics.keys()})"
        )
    current_val = trainer.callback_metrics[self.monitor].item()

    is_better = self._best_value is None or (
        (self.mode == "min" and current_val < self._best_value)
        or (self.mode == "max" and current_val > self._best_value)
    )
    if is_better:
        self._best_value = current_val
        best_path = os.path.join(self.dirpath, "best.ckpt")
        trainer.save_checkpoint(best_path)

        # Previously this entire function was decorated with @rank_zero_only, but this resulted
        # in NCCL timeout because all ranks must call trainer.save_checkpoint. Log only on rank 0.
        if trainer.is_global_zero:
            logger.info(
                "saved best checkpoint to %s (%s=%s)",
                best_path,
                self.monitor,
                current_val,
            )

    # Save last.ckpt after best.ckpt so that auto-resuming from last.ckpt is always
    # correct.
    # - If we crash after writing best.ckpt, we will auto-resume from previous epoch
    #   and write the same best.ckpt again.
    # - If we crash after writing last.ckpt, we will restore the correct
    #   self._best_value and won't override the checkpoint incorrectly.
    last_path = os.path.join(self.dirpath, "last.ckpt")
    trainer.save_checkpoint(last_path)
    if trainer.is_global_zero:
        logger.info("saved checkpoint to %s", last_path)

ManagedBestLastCheckpoint

Bases: BestLastCheckpoint

BestLastCheckpoint that resolves dirpath from trainer.default_root_dir.

Use with project management: the CLI sets trainer.default_root_dir to the project directory ({MANAGEMENT_DIR}/{PROJECT_NAME}/{RUN_NAME}), and this callback picks it up at setup time.

Source code in rslearn/train/callbacks/checkpointing.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
class ManagedBestLastCheckpoint(BestLastCheckpoint):
    """BestLastCheckpoint that resolves dirpath from trainer.default_root_dir.

    Use with project management: the CLI sets trainer.default_root_dir to the project
    directory ({MANAGEMENT_DIR}/{PROJECT_NAME}/{RUN_NAME}), and this callback picks it
    up at setup time.
    """

    def __init__(
        self, monitor: str = "val_loss", mode: Literal["min", "max"] = "min"
    ) -> None:
        """Create a new ManagedBestLastCheckpoint.

        Args:
            monitor: metric key to monitor for best checkpoint.
            mode: "min" if lower is better, "max" if higher is better.
        """
        super().__init__(dirpath="", monitor=monitor, mode=mode)

    def setup(
        self, trainer: Trainer, pl_module: LightningModule, stage: str | None = None
    ) -> None:
        """Resolve dirpath from trainer.default_root_dir."""
        self.dirpath = trainer.default_root_dir
        if trainer.is_global_zero:
            logger.info("ManagedBestLastCheckpoint using dirpath=%s", self.dirpath)

setup 

setup(trainer: Trainer, pl_module: LightningModule, stage: str | None = None) -> None

Resolve dirpath from trainer.default_root_dir.

Source code in rslearn/train/callbacks/checkpointing.py 
105
106
107
108
109
110
111
def setup(
    self, trainer: Trainer, pl_module: LightningModule, stage: str | None = None
) -> None:
    """Resolve dirpath from trainer.default_root_dir."""
    self.dirpath = trainer.default_root_dir
    if trainer.is_global_zero:
        logger.info("ManagedBestLastCheckpoint using dirpath=%s", self.dirpath)