Skip to content
This repository was archived by the owner on Jun 3, 2025. It is now read-only.

[BugFix][Torchvision] update optimizer state dict before transfer learning #1358

Merged
merged 5 commits into from
Feb 2, 2023
Merged

[BugFix][Torchvision] update optimizer state dict before transfer learning #1358

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
de64598
Add: an `_update_checkpoint_optimizer(...)` for deleting mismatching …
rahul-tuli Feb 1, 2023
4eb372a
Remove: _update_checkpoint_optimizer in favor of loading in the optim…
rahul-tuli Feb 2, 2023
31e3172
Remove: un-needed imports
rahul-tuli Feb 2, 2023
91f020f
Address review comments
rahul-tuli Feb 2, 2023
960db50
Style
rahul-tuli Feb 2, 2023
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
9 changes: 8 additions & 1 deletion src/sparseml/pytorch/torchvision/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -533,7 +533,14 @@ def collate_fn(batch):
# load params
if checkpoint is not None:
if "optimizer" in checkpoint and not args.test_only:
optimizer.load_state_dict(checkpoint["optimizer"])
if args.resume:
optimizer.load_state_dict(checkpoint["optimizer"])
else:
warnings.warn(
"Optimizer state dict not loaded from checkpoint. Unless run is "
"resumed with the --resume arg, the optimizer will start from a "
"fresh state"
)
if model_ema and "model_ema" in checkpoint:
model_ema.load_state_dict(checkpoint["model_ema"])
if scaler and "scaler" in checkpoint:
Expand Down