Skip to content
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

SAM configuration check fixes #288

Merged
merged 5 commits into from
Sep 10, 2024
Merged

SAM configuration check fixes #288

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6f668a9
fix: correcting mutual exclusivity check
laserkelvin Sep 10, 2024
8134387
fix: correcting float check for epoch skip setting
laserkelvin Sep 10, 2024
dc3021e
fix: correcting is_integer check to float instance checking
laserkelvin Sep 10, 2024
60a04e7
refactor: configuring logger to report when SAM is expected to start
laserkelvin Sep 10, 2024
ff0a415
refactor: adding step configuration to SAM example
laserkelvin Sep 10, 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
6 changes: 4 additions & 2 deletions examples/callbacks/SAM.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@
},
)

# run a quick training loop
trainer = pl.Trainer(fast_dev_run=1000, callbacks=[SAM()])
# run a quick training loop, skipping the first five steps
trainer = pl.Trainer(
fast_dev_run=100, callbacks=[SAM(skip_step_count=5, logging=True, log_level="INFO")]
)
trainer.fit(task, datamodule=dm)
10 changes: 7 additions & 3 deletions matsciml/lightning/callbacks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -776,12 +776,12 @@ def __init__(
super().__init__()
self.rho = rho
self.adaptive = adaptive
if skip_epoch_count and skip_epoch_count:
if skip_step_count and skip_epoch_count:
raise ValueError(
"`skip_epoch_count` and `skip_step_count` are mutually exclusive for SAM."
)
self.skip_step_count = skip_step_count
if skip_epoch_count and not skip_epoch_count.is_integer():
if skip_epoch_count and isinstance(skip_epoch_count, float):
assert (
0 < skip_epoch_count < 1.0
), "Decimal `skip_epoch_count` passed not within [0,1]."
Expand All @@ -808,8 +808,12 @@ def on_fit_start(
self.max_steps = train_len * self.max_epochs
# if a fractional epoch skip is specified, convert it to
# an integer count for easier comparison
if self.skip_epoch_count and not self.skip_epoch_count.is_integer():
if self.skip_epoch_count and isinstance(self.skip_epoch_count, float):
self.skip_epoch_count = int(self.max_epochs * self.skip_epoch_count)
if self.logger:
self.logger.info(
f"Fractional epoch skip - will start SAM from epoch {self.skip_epoch_count}, max epochs {self.max_epochs}."
)
# add floating point epsilon for later use
self.epsilon = torch.tensor(
[torch.finfo(pl_module.dtype).eps],
Expand Down
Loading