Skip to content

Commit

Permalink
[Major] simplify learning rate finder and use log-mean of default and…
Browse files Browse the repository at this point in the history
… smooth suggestion (#1630)

* simplify learning rate finder

* set learning rate to logarithmic mean of smooth and default.

* move logic to util
  • Loading branch information
ourownstory authored Aug 14, 2024
1 parent 982e152 commit 8330bab
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 10 deletions.
4 changes: 2 additions & 2 deletions neuralprophet/configure.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ def set_lr_finder_args(self, dataset_size, num_batches):
Set the lr_finder_args.
This is the range of learning rates to test.
"""
num_training = 150 + int(np.log10(100 + dataset_size) * 25)
num_training = 100 + int(np.log10(dataset_size) * 20)
if num_batches < num_training:
log.warning(
f"Learning rate finder: The number of batches ({num_batches}) is too small than the required number \
Expand All @@ -217,7 +217,7 @@ def set_lr_finder_args(self, dataset_size, num_batches):
# num_training = num_batches
self.lr_finder_args.update(
{
"min_lr": 1e-6,
"min_lr": 1e-7,
"max_lr": 10,
"num_training": num_training,
"early_stop_threshold": None,
Expand Down
8 changes: 3 additions & 5 deletions neuralprophet/forecaster.py
Original file line number Diff line number Diff line change
Expand Up @@ -2805,13 +2805,12 @@ def _train(
lr_finder = tuner.lr_find(
model=self.model,
train_dataloaders=train_loader,
val_dataloaders=val_loader,
# val_dataloaders=val_loader, # not be used, but may lead to Lightning bug if not provided
**self.config_train.lr_finder_args,
)
# Estimate the optimal learning rate from the loss curve
assert lr_finder is not None
_, _, lr_suggestion = utils.smooth_loss_and_suggest(lr_finder.results)
self.model.learning_rate = lr_suggestion
_, _, self.model.learning_rate = utils.smooth_loss_and_suggest(lr_finder)
start = time.time()
self.trainer.fit(
self.model,
Expand All @@ -2832,8 +2831,7 @@ def _train(
)
assert lr_finder is not None
# Estimate the optimal learning rate from the loss curve
_, _, lr_suggestion = utils.smooth_loss_and_suggest(lr_finder.results)
self.model.learning_rate = lr_suggestion
_, _, self.model.learning_rate = utils.smooth_loss_and_suggest(lr_finder)
start = time.time()
self.trainer.fit(
self.model,
Expand Down
18 changes: 15 additions & 3 deletions neuralprophet/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -751,7 +751,7 @@ def set_log_level(log_level: str = "INFO", include_handlers: bool = False):
set_logger_level(logging.getLogger("NP"), log_level, include_handlers)


def smooth_loss_and_suggest(lr_finder_results, window=10):
def smooth_loss_and_suggest(lr_finder, window=10):
"""
Smooth loss using a Hamming filter.
Expand All @@ -769,10 +769,12 @@ def smooth_loss_and_suggest(lr_finder_results, window=10):
suggested_lr: float
Suggested learning rate based on gradient
"""
lr_finder_results = lr_finder.results
lr = lr_finder_results["lr"]
loss = lr_finder_results["loss"]
# Derive window size from num lr searches, ensure window is divisible by 2
half_window = math.ceil(round(len(loss) * 0.1) / 2)
# half_window = math.ceil(round(len(loss) * 0.1) / 2)
half_window = math.ceil(window / 2)
# Pad sequence and initialialize hamming filter
loss = np.pad(np.array(loss), pad_width=half_window, mode="edge")
window = np.hamming(half_window * 2)
Expand All @@ -798,7 +800,17 @@ def smooth_loss_and_suggest(lr_finder_results, window=10):
"samples or manually set the learning rate."
)
raise
return (loss, lr, suggestion)
suggestion_default = lr_finder.suggestion(skip_begin=10, skip_end=3)
if suggestion is not None and suggestion_default is not None:
log_suggestion_smooth = np.log(suggestion)
log_suggestion_default = np.log(suggestion_default)
lr_suggestion = np.exp((log_suggestion_smooth + log_suggestion_default) / 2)
elif suggestion is None and suggestion_default is None:
log.error("Automatic learning rate test failed. Please set manually the learning rate.")
raise
else:
lr_suggestion = suggestion if suggestion is not None else suggestion_default
return (loss, lr, lr_suggestion)


def _smooth_loss(loss, beta=0.9):
Expand Down

0 comments on commit 8330bab

Please sign in to comment.