Skip to content

Added workaround for CyclicLR saving #1683

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
anautsch merged 1 commit into from
Nov 3, 2022
Merged

Added workaround for CyclicLR saving #1683

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
33 changes: 33 additions & 0 deletions speechbrain/utils/_workarounds.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
"""This module implements some workarounds for dependencies

Authors
* Aku Rouhe 2022
"""
import torch
import weakref
import warnings

WEAKREF_MARKER = "WEAKREF"


def _cycliclrsaver(obj, path):
state_dict = obj.state_dict()
if state_dict.get("_scale_fn_ref") is not None:
state_dict["_scale_fn_ref"] = WEAKREF_MARKER
torch.save(state_dict, path)


def _cycliclrloader(obj, path, end_of_epoch, device=None):
del end_of_epoch # Unused
state_dict = torch.load(path, map_location=device)
if state_dict.get("_scale_fn_ref") == WEAKREF_MARKER:
if not isinstance(obj._scale_fn_ref, weakref.WeakMethod):
MSG = "Loading CyclicLR scheduler and the _scale_ref_fn did not exist in instance."
MSG += " You did not construct it with the same parameters it was created!"
MSG += " Looks like you changed the scale function!"
MSG += " If this was not intentional, the scheduler might not work correctly."
warnings.warn(MSG)
try:
obj.load_state_dict(torch.load(path, map_location=device), strict=True)
except TypeError:
obj.load_state_dict(torch.load(path, map_location=device))
5 changes: 5 additions & 0 deletions speechbrain/utils/checkpoints.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import shutil
import logging
import warnings
import speechbrain.utils._workarounds as __wa

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -188,6 +189,10 @@ def _load_spm(obj, path, device=None):
# SentencePiece not loaded, fine!
pass

# Add workarounds:
DEFAULT_SAVE_HOOKS[torch.optim.lr_scheduler.CyclicLR] = __wa._cycliclrsaver
DEFAULT_LOAD_HOOKS[torch.optim.lr_scheduler.CyclicLR] = __wa._cycliclrloader


def mark_as_saver(method):
"""Method decorator which marks given method as the checkpoint saving hook.
Expand Down