Skip to content

deeplabcut.pose_estimation_pytorch.runners.snapshots

Code to handle storing models.

Classes:

Name Description
TorchSnapshotManager

Class handling model checkpoint I/O.

TorchSnapshotManager dataclass

Class handling model checkpoint I/O.

Attributes:

Name Type Description
snapshot_prefix str

The prefix to use when saving snapshots.
model_folder Path

The path to the directory where model snapshots should be stored.
key_metric str | None

If defined, the metric is used to save the best model. Otherwise no best model is used.
key_metric_asc bool

Whether the key metric is ascending (larger values are better).
max_snapshots int

The maximum number of snapshots to store for the training run. This does not include the best model (e.g., setting max_snapshots=5 will mean that the 5 latest models will be kept, plus the best model)
save_epochs int

The number of epochs between each model save
save_optimizer_state bool

Whether to store the optimizer state. This makes snapshots much heavier, but allows to resume training as if it was never stopped.

Examples:

Storing snapshots while training

model: nn.Module loader = DLCLoader(...) snapshot_manager = TorchSnapshotManager( "snapshot", loader.model_folder, key_metric="test.mAP", ) ... for epoch in range(num_epochs): train_epoch(model, data) snapshot_manager.update({ "metadata": { "metrics": {"mAP": ...} }, "model": model.state_dict(), "optimizer": optimizer.state_dict() })

Methods:

Name Description
best

Returns: the path to the best snapshot, if it exists
last

Returns: path to the last snapshot that was saved, if any snapshot exists
snapshot_path

Args:
snapshots

Args:
update

Saves the model state dict if the epoch is one that requires a save.
Source code in deeplabcut/pose_estimation_pytorch/runners/snapshots.py 
 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
 84
 85
 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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
@dataclass
class TorchSnapshotManager:
    """Class handling model checkpoint I/O.

    Attributes:
        snapshot_prefix: The prefix to use when saving snapshots.
        model_folder: The path to the directory where model snapshots should be stored.
        key_metric: If defined, the metric is used to save the best model. Otherwise no
            best model is used.
        key_metric_asc: Whether the key metric is ascending (larger values are better).
        max_snapshots: The maximum number of snapshots to store for the training run.
            This does not include the best model (e.g., setting max_snapshots=5 will
            mean that the 5 latest models will be kept, plus the best model)
        save_epochs: The number of epochs between each model save
        save_optimizer_state: Whether to store the optimizer state. This makes snapshots
            much heavier, but allows to resume training as if it was never stopped.

    Examples:
        # Storing snapshots while training
        model: nn.Module
        loader = DLCLoader(...)
        snapshot_manager = TorchSnapshotManager(
            "snapshot",
            loader.model_folder,
            key_metric="test.mAP",
        )
        ...
        for epoch in range(num_epochs):
            train_epoch(model, data)
            snapshot_manager.update({
                "metadata": {
                    "metrics": {"mAP": ...}
                },
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict()
            })
    """

    snapshot_prefix: str
    model_folder: Path
    key_metric: str | None = None
    key_metric_asc: bool = True
    max_snapshots: int = 5
    save_epochs: int = 25
    save_optimizer_state: bool = False
    _best_model_epochs: int = -1
    _best_metric: float | None = None
    _key: str = field(init=False)

    def __post_init__(self):
        assert self.max_snapshots > 0, "max_snapshots must be a positive integer"
        self._key = f"metrics/{self.key_metric}"

    def update(self, epoch: int, state_dict: dict, last: bool = False) -> None:
        """Saves the model state dict if the epoch is one that requires a save.

        Args:
            epoch: the number of epochs the model was trained for
            state_dict: the state dict to store
            last: whether this is the last epoch in the training run, which forces a
                model save no matter the epoch number

        Returns:
            the path to the saved snapshot if one
        """
        metrics = state_dict["metadata"]["metrics"]
        if (
            self._key in metrics
            and not np.isnan(metrics[self._key])
            and (
                self._best_metric is None
                or (self.key_metric_asc and self._best_metric < metrics[self._key])
                or (not self.key_metric_asc and self._best_metric > metrics[self._key])
            )
        ):
            current_best = self.best()
            self._best_metric = metrics[self._key]

            # Save the new best model
            save_path = self.snapshot_path(epoch, best=True)
            parsed_state_dict = {k: v for k, v in state_dict.items() if self.save_optimizer_state or k != "optimizer"}
            torch.save(parsed_state_dict, save_path)

            # Handle previous best model
            if current_best is not None:
                if current_best.epochs % self.save_epochs == 0:
                    new_name = self.snapshot_path(epoch=current_best.epochs)
                    current_best.path.rename(new_name)
                else:
                    current_best.path.unlink(missing_ok=False)
        elif last or epoch % self.save_epochs == 0:
            # Save regular snapshot if needed
            save_path = self.snapshot_path(epoch=epoch)
            parsed_state_dict = {k: v for k, v in state_dict.items() if self.save_optimizer_state or k != "optimizer"}
            torch.save(parsed_state_dict, save_path)

        # Clean up old snapshots if needed
        existing_snapshots = [s for s in self.snapshots() if not s.best]
        if len(existing_snapshots) >= self.max_snapshots:
            num_to_delete = len(existing_snapshots) - self.max_snapshots
            to_delete = existing_snapshots[:num_to_delete]
            for snapshot in to_delete:
                snapshot.path.unlink(missing_ok=False)

    def best(self) -> Snapshot | None:
        """Returns: the path to the best snapshot, if it exists"""
        snapshots = self.snapshots()
        best_snapshots = [s for s in snapshots if s.best]
        if len(best_snapshots) == 0:
            return None

        if len(best_snapshots) > 1:
            warnings.warn(
                f"TorchSnapshotManager.best(): found multiple best snapshots ("
                f"{best_snapshots}), returning the last one.",
                stacklevel=2,
            )

        best_snapshot = best_snapshots[-1]
        return best_snapshot

    def last(self) -> Snapshot | None:
        """Returns: path to the last snapshot that was saved, if any snapshot exists"""
        snapshots = self.snapshots(best_in_last=False)
        if len(snapshots) == 0:
            return None
        return snapshots[-1]

    def snapshots(self, best_in_last: bool = True) -> list[Snapshot]:
        """
        Args:
            best_in_last: Whether to place the snapshot with the best performance in the
                last position in the list, even if it wasn't the last epoch.

        Returns:
            The snapshots for a training run, sorted by the number of epochs they were
            trained for. If ``best_in_last=True`` and a best snapshot exists, it will be
            the last one in the list.
        """
        return list_snapshots(self.model_folder, self.snapshot_prefix, best_in_last=best_in_last)

    def snapshot_path(self, epoch: int, best: bool = False) -> Path:
        """
        Args:
            epoch: the number of epochs for which a snapshot was trained
            best: whether this is the best performing model for the training run

        Returns:
            the path where the model should be stored
        """
        uid = f"{epoch:03}"
        if best:
            uid = f"best-{uid}"
        return self.model_folder / f"{self.snapshot_prefix}-{uid}.pt"

best 

best() -> Snapshot | None

Returns: the path to the best snapshot, if it exists

Source code in deeplabcut/pose_estimation_pytorch/runners/snapshots.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def best(self) -> Snapshot | None:
    """Returns: the path to the best snapshot, if it exists"""
    snapshots = self.snapshots()
    best_snapshots = [s for s in snapshots if s.best]
    if len(best_snapshots) == 0:
        return None

    if len(best_snapshots) > 1:
        warnings.warn(
            f"TorchSnapshotManager.best(): found multiple best snapshots ("
            f"{best_snapshots}), returning the last one.",
            stacklevel=2,
        )

    best_snapshot = best_snapshots[-1]
    return best_snapshot

last 

last() -> Snapshot | None

Returns: path to the last snapshot that was saved, if any snapshot exists

Source code in deeplabcut/pose_estimation_pytorch/runners/snapshots.py 
146
147
148
149
150
151
def last(self) -> Snapshot | None:
    """Returns: path to the last snapshot that was saved, if any snapshot exists"""
    snapshots = self.snapshots(best_in_last=False)
    if len(snapshots) == 0:
        return None
    return snapshots[-1]

snapshot_path 

snapshot_path(epoch: int, best: bool = False) -> Path

Parameters:

Name Type Description Default

epoch

int

the number of epochs for which a snapshot was trained

 required

best

bool

whether this is the best performing model for the training run

 False

Returns:

Type Description
Path

the path where the model should be stored
Source code in deeplabcut/pose_estimation_pytorch/runners/snapshots.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
def snapshot_path(self, epoch: int, best: bool = False) -> Path:
    """
    Args:
        epoch: the number of epochs for which a snapshot was trained
        best: whether this is the best performing model for the training run

    Returns:
        the path where the model should be stored
    """
    uid = f"{epoch:03}"
    if best:
        uid = f"best-{uid}"
    return self.model_folder / f"{self.snapshot_prefix}-{uid}.pt"

snapshots 

snapshots(best_in_last: bool = True) -> list[Snapshot]

Parameters:

Name Type Description Default

best_in_last

bool

Whether to place the snapshot with the best performance in the last position in the list, even if it wasn't the last epoch.

 True

Returns:

Type Description
list[Snapshot]

The snapshots for a training run, sorted by the number of epochs they were trained for. If best_in_last=True and a best snapshot exists, it will be the last one in the list.
Source code in deeplabcut/pose_estimation_pytorch/runners/snapshots.py 
153
154
155
156
157
158
159
160
161
162
163
164
def snapshots(self, best_in_last: bool = True) -> list[Snapshot]:
    """
    Args:
        best_in_last: Whether to place the snapshot with the best performance in the
            last position in the list, even if it wasn't the last epoch.

    Returns:
        The snapshots for a training run, sorted by the number of epochs they were
        trained for. If ``best_in_last=True`` and a best snapshot exists, it will be
        the last one in the list.
    """
    return list_snapshots(self.model_folder, self.snapshot_prefix, best_in_last=best_in_last)

update 

update(epoch: int, state_dict: dict, last: bool = False) -> None

Saves the model state dict if the epoch is one that requires a save.

Parameters:

Name Type Description Default

epoch

int

the number of epochs the model was trained for

 required

state_dict

dict

the state dict to store

 required

last

bool

whether this is the last epoch in the training run, which forces a model save no matter the epoch number

 False

Returns:

Type Description
None

the path to the saved snapshot if one
Source code in deeplabcut/pose_estimation_pytorch/runners/snapshots.py 
 78
 79
 80
 81
 82
 83
 84
 85
 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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def update(self, epoch: int, state_dict: dict, last: bool = False) -> None:
    """Saves the model state dict if the epoch is one that requires a save.

    Args:
        epoch: the number of epochs the model was trained for
        state_dict: the state dict to store
        last: whether this is the last epoch in the training run, which forces a
            model save no matter the epoch number

    Returns:
        the path to the saved snapshot if one
    """
    metrics = state_dict["metadata"]["metrics"]
    if (
        self._key in metrics
        and not np.isnan(metrics[self._key])
        and (
            self._best_metric is None
            or (self.key_metric_asc and self._best_metric < metrics[self._key])
            or (not self.key_metric_asc and self._best_metric > metrics[self._key])
        )
    ):
        current_best = self.best()
        self._best_metric = metrics[self._key]

        # Save the new best model
        save_path = self.snapshot_path(epoch, best=True)
        parsed_state_dict = {k: v for k, v in state_dict.items() if self.save_optimizer_state or k != "optimizer"}
        torch.save(parsed_state_dict, save_path)

        # Handle previous best model
        if current_best is not None:
            if current_best.epochs % self.save_epochs == 0:
                new_name = self.snapshot_path(epoch=current_best.epochs)
                current_best.path.rename(new_name)
            else:
                current_best.path.unlink(missing_ok=False)
    elif last or epoch % self.save_epochs == 0:
        # Save regular snapshot if needed
        save_path = self.snapshot_path(epoch=epoch)
        parsed_state_dict = {k: v for k, v in state_dict.items() if self.save_optimizer_state or k != "optimizer"}
        torch.save(parsed_state_dict, save_path)

    # Clean up old snapshots if needed
    existing_snapshots = [s for s in self.snapshots() if not s.best]
    if len(existing_snapshots) >= self.max_snapshots:
        num_to_delete = len(existing_snapshots) - self.max_snapshots
        to_delete = existing_snapshots[:num_to_delete]
        for snapshot in to_delete:
            snapshot.path.unlink(missing_ok=False)