Repository: ludwig-ai/ludwig
Pull request: Fix LR reduce on plateau interaction with base LR decay
Fixes #3049.

Problem statement:
Learning rate reduction not functional
**Describe the bug**
Despite the respective settings in the `trainer` section, the learning rate does not decay when a plateau is reached.

**To Reproduce**

I am using ludwig for image classification with the following .yaml
```
input_features:
- name: image_path
  type: image
  encoder: resnet
  resnet_size: 101
  num_filters: 16
  kernel_size: 3
  conv_stride: 1
  batch_norm_momentum: 0.9
  batch_norm_epsilon: 0.001
  num_fc_layers: 1
  output_size: 256
  use_bias: true
  bias_initializer: zeros
  norm: null
  norm_params: null
  activation: relu
  dropout: 0

output_features:
 - name: class
   type: category

trainer:
  epochs: 50
  batch_size: 16
  early_stop: 8
  learning_rate: 0.001
  reduce_learning_rate_on_plateau: 4
  reduce_learning_rate_on_plateau_patience: 4
  reduce_learning_rate_on_plateau_rate: 0.5
```
I train the model distributed with horovod in a ludwig stock docker container:
```
horovodrun -np 3 ludwig train --config model_definition.yaml --dataset annotations.csv --experiment_name resnet101
```

In this particular case the learning rate the model describes is even `0.003`. An Excerpt from `training_progress.json`:

```
    "best_increase_batch_size_eval_metric": Infinity,
    "checkpoint_number": 34,
    "epoch": 34,
    "last_improvement_steps": 6512,
    "last_increase_batch_size": 0,
    "last_increase_batch_size_eval_metric_improvement": 0,
    "last_increase_batch_size_steps": 0,
    "last_learning_rate_reduction": 0,
    "last_learning_rate_reduction_steps": 0,
    "learning_rate": 0.003,
    "num_increases_batch_size": 0,
    "num_reductions_learning_rate": 0,
    "steps": 27676,
```

**Expected behavior**
The learning rate should be reduced after 4 epochs on no further loss improvement/

**Environment (please complete the following information):**

- OS: Linux Docker environment
- Version: Docker ludwigai/ludwig-ray-gpu:master
- Python version 3.8.5
- Ludwig version 0.7.beta

**Additional context**
this does not seem to be ludwig version specific. I experienced a similar behavior with ludwig 0.5.5

Thanks so much for some feedback and keep up the great work.

Hints from issue discussion:
Hey @tboo, thanks for reporting this issue. We recently made some big changes to learning rate reduction, so definitely want to make sure it's working correctly! 

I was able to repro the issue you're seeing in your environment and will take some time to dig into it tomorrow.

Commit message: Fixed progress tracker

Patch to review:
diff --git a/ludwig/modules/lr_scheduler.py b/ludwig/modules/lr_scheduler.py
index 5a1e03598af..cc1a2b06747 100644
--- a/ludwig/modules/lr_scheduler.py
+++ b/ludwig/modules/lr_scheduler.py
@@ -25,9 +25,26 @@ def step(self, metrics):
 
         return super().step(metrics)
 
+    @property
+    def num_reduce_lr(self) -> int:
+        return self._num_reduce_lr
+
     def _reduce_lr(self, epoch):
-        super()._reduce_lr(epoch)
         self._num_reduce_lr += 1
+        self.apply_lr(epoch)
+
+    def apply_lr(self, epoch=None):
+        if self._num_reduce_lr == 0:
+            return
+
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            old_lr = float(param_group["lr"])
+            new_lr = max(old_lr * math.pow(self.factor, self._num_reduce_lr), self.min_lrs[i])
+            if old_lr - new_lr > self.eps:
+                param_group["lr"] = new_lr
+                if self.verbose:
+                    epoch_str = ("%.2f" if isinstance(epoch, float) else "%.5d") % epoch
+                    print("Epoch {}: reducing learning rate" " of group {} to {:.4e}.".format(epoch_str, i, new_lr))
 
 
 class LRScheduler:
@@ -67,6 +84,12 @@ def reset(self, steps_per_checkpoint: int, total_steps: int):
     def step(self):
         self._train_scheduler.step()
 
+        if self._eval_scheduler is not None:
+            # We apply this scheduler every eval step, not train step, so we don't want to call step() here.
+            # However, we need to re-apply the LR reduction to the LR from the train scheduler, as the first scheduler
+            # resets the LR back to the base LR.
+            self._eval_scheduler.apply_lr()
+
     def eval_step(self, progress_tracker: ProgressTracker, validation_field: str):
         if self._eval_scheduler is None:
             # No reduce on plateau
@@ -83,8 +106,18 @@ def eval_step(self, progress_tracker: ProgressTracker, validation_field: str):
         last_metric: TrainerMetric = split_metrics[validation_field][validation_metric][-1]
         last_metric_value = last_metric[-1]
 
+        prev_num_reductions = self._eval_scheduler.num_reduce_lr
         self._eval_scheduler.step(last_metric_value)
 
+        num_reductions = self._eval_scheduler.num_reduce_lr
+        if num_reductions > prev_num_reductions:
+            # LR reduction -> update progress tracker
+            progress_tracker.last_learning_rate_reduction_steps = progress_tracker.steps
+            progress_tracker.last_learning_rate_reduction = 0
+            progress_tracker.num_reductions_learning_rate += 1
+        else:
+            progress_tracker.steps - progress_tracker.last_learning_rate_reduction_steps
+
     def state_dict(self) -> Dict[str, Any]:
         return {
             "train_scheduler_state": self._train_scheduler.state_dict(),
diff --git a/tests/ludwig/modules/test_lr_scheduler.py b/tests/ludwig/modules/test_lr_scheduler.py
index 10270d6b273..a87448c8639 100644
--- a/tests/ludwig/modules/test_lr_scheduler.py
+++ b/tests/ludwig/modules/test_lr_scheduler.py
@@ -1,3 +1,4 @@
+import numpy as np
 from torch.optim import SGD
 
 from ludwig.features.number_feature import NumberInputFeature, NumberOutputFeature
@@ -72,7 +73,13 @@ def test_lr_scheduler_reduce_on_plateau():
     )
 
     optimizer = SGD(module.parameters(), lr=base_lr)
-    config = LRSchedulerConfig(warmup_evaluations=0, reduce_on_plateau=reduce_limit)
+    config = LRSchedulerConfig(
+        warmup_evaluations=0,
+        decay=None,
+        reduce_on_plateau=reduce_limit,
+        reduce_on_plateau_patience=10,
+        reduce_on_plateau_rate=0.1,
+    )
     scheduler = LRScheduler(config=config, optimizer=optimizer)
 
     progress_tracker = get_new_progress_tracker(
@@ -89,6 +96,11 @@ def test_lr_scheduler_reduce_on_plateau():
     steps_to_plateau = 5
     loss = 10.0
     for epoch in range(total_eval_steps):
+        for i in range(100):
+            # Simulate batch-wise steps. If we make a mistake, then this will reset
+            # the learning rate.
+            scheduler.step()
+
         steps_to_plateau -= 1
         if steps_to_plateau > 0:
             loss -= 0.1
@@ -106,6 +118,9 @@ def test_lr_scheduler_reduce_on_plateau():
 
     assert num_reductions == reduce_limit
 
+    # 3 reductions that multiply by 0.1 each time
+    assert np.isclose(lr, 0.001)
+
 
 def test_lr_scheduler_save_load():
     steps_per_checkpoint = 10
@@ -139,7 +154,6 @@ def test_lr_scheduler_save_load():
 
     optimizer_state = optimizer.state_dict()
     scheduler_state = scheduler.state_dict()
-    print(scheduler_state)
 
     optimizer2 = SGD(module.parameters(), lr=base_lr)
     scheduler2 = LRScheduler(config=config, optimizer=optimizer2)


Code review comments:
diff_hunk: @@ -83,8 +106,18 @@ def eval_step(self, progress_tracker: ProgressTracker, validation_field: str):
         last_metric: TrainerMetric = split_metrics[validation_field][validation_metric][-1]
         last_metric_value = last_metric[-1]
 
+        prev_num_reductions = self._eval_scheduler.num_reduce_lr
         self._eval_scheduler.step(last_metric_value)
 
+        num_reductions = self._eval_scheduler.num_reduce_lr
+        if num_reductions > prev_num_reductions:
+            # LR reduction -> update progress tracker
+            progress_tracker.last_learning_rate_reduction_steps = progress_tracker.steps
+            progress_tracker.last_learning_rate_reduction = 0
+            progress_tracker.num_reductions_learning_rate += 1
+        else:
+            progress_tracker.steps - progress_tracker.last_learning_rate_reduction_steps
path: ludwig/modules/lr_scheduler.py
text: @user1:
Should this be: `progress_tracker.last_learning_rate_reduction_steps = progress_tracker.steps - progress_tracker.last_learning_rate_reduction_steps`?

@author:
Ah, good catch!