Visitar URL original
Fix: correct the early return error when save_epochs=1 and optimize the clean up way by xiu-cs · Pull Request #2793 · DeepLabCut/DeepLabCut · GitHub
Skip to content

Fix: correct the early return error when save_epochs=1 and optimize the clean up way #2793

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
n-poulsen merged 4 commits into from
Nov 25, 2024
Merged

Fix: correct the early return error when save_epochs=1 and optimize the clean up way #2793

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
598d02b
Fix: correct the early return error when save_epochs=1 and delelte th…
xiu-cs Nov 21, 2024
76e10b4
Update deeplabcut/pose_estimation_pytorch/runners/snapshots.py
xiu-cs Nov 22, 2024
0b64247
Merge branch 'pytorch_dlc' into ti/snapshots
xiu-cs Nov 22, 2024
f670407
black -> snapshots.py
xiu-cs Nov 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 18 additions & 16 deletions deeplabcut/pose_estimation_pytorch/runners/snapshots.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,41 +113,43 @@ def update(self, epoch: int, state_dict: dict, last: bool = False) -> None:
):
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
k: v
Comment thread
xiu-cs marked this conversation as resolved.
Resolved
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:
# rename if the current best should have been saved, otherwise delete
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)
return

if not (last or epoch % self.save_epochs == 0):
return

else:
# Save regular snapshot if needed
should_save = last or epoch % self.save_epochs == 0
if should_save:
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)
Comment thread
xiu-cs marked this conversation as resolved.
Resolved

# 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 = 1 + 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)

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)

def best(self) -> Snapshot | None:
"""Returns: the path to the best snapshot, if it exists"""
snapshots = self.snapshots()
Expand Down