Repository: pytorch/ignite

Issue:
Improve create_lr_scheduler_with_warmup
Actual behavior of `create_lr_scheduler_with_warmup` is a bit incoherent and varies upon pytorch version (1.3.0 vs nightly).

Actual problems:
- warmup duration is +1 event longer than expected on torch lr schedulers
- learning rate after warmup for torch lr schedulers should follow optimizer's initial lr
- first learning rate after warmup is skipped in discontinuous case

PyTorch v1.3.0
![image](https://user-images.githubusercontent.com/2459423/67299264-b5bc8180-f4ec-11e9-985d-a41769b702e5.png)

![image](https://user-images.githubusercontent.com/2459423/67299291-c10fad00-f4ec-11e9-9954-8dd1cc1ffde4.png)

PyTorch v1.4.dev
![image](https://user-images.githubusercontent.com/2459423/67299551-1946af00-f4ed-11e9-9b69-2b226c040330.png)

![image](https://user-images.githubusercontent.com/2459423/67299566-1ea3f980-f4ed-11e9-8aa1-1c1a8f1bc536.png)


## Solution

- same warmup duration for all cases
- add additional event in discontinuos lr scheduling
- same behaviour for PyTorch v1.3.0 and nightly

![image](https://user-images.githubusercontent.com/2459423/67299842-73e00b00-f4ed-11e9-9e7f-560ea428f9fa.png)

![image](https://user-images.githubusercontent.com/2459423/67299860-793d5580-f4ed-11e9-9407-7d58cf040c9c.png)

 
### Code to reproduce
```python
import numpy as np
import matplotlib.pylab as plt
%matplotlib inline

import torch
import ignite
from ignite.engine import Engine, Events
from ignite.contrib.handlers import LinearCyclicalScheduler, create_lr_scheduler_with_warmup
from ignite.contrib.handlers.param_scheduler import ParamScheduler

configs = [
  {
    "name": "create_lr_scheduler_with_warmup with torch LR scheduler",
    "description": "Optimizer init lr={}, warmup_end_value={}, warmup duration={}",
    "init_lr": 0.01,
    "warmup_end_value": 0.03,
    "lr_scheduler_cls": torch.optim.lr_scheduler.StepLR,
    "lr_scheduler_kwargs": {"step_size": 3, "gamma": 0.5},
  },
  {
    "name": "create_lr_scheduler_with_warmup with torch LR scheduler",
    "description": "Optimizer init lr={}, warmup_end_value={}, warmup duration={}",
    "init_lr": 0.01,
    "warmup_end_value": 0.01,
    "lr_scheduler_cls": torch.optim.lr_scheduler.StepLR,
    "lr_scheduler_kwargs": {"step_size": 3, "gamma": 0.5},    
  },
  {
    "name": "create_lr_scheduler_with_warmup with ignite param scheduler",
    "description": "param scheduler's init lr={}, warmup_end_value={}, warmup duration={}",
    "init_lr": 0.0,
    "warmup_end_value": 0.03,
    "lr_scheduler_cls": LinearCyclicalScheduler,
    "lr_scheduler_kwargs": {"param_name": "lr", "start_value": 0.01, "end_value": 0.0, "cycle_size": 12},    
  },
  {
    "name": "create_lr_scheduler_with_warmup with ignite param scheduler",
    "description": "param scheduler's init lr={}, warmup_end_value={}, warmup duration={}",
    "init_lr": 0.0,
    "warmup_end_value": 0.01,
    "lr_scheduler_cls": LinearCyclicalScheduler,
    "lr_scheduler_kwargs": {"param_name": "lr", "start_value": 0.01, "end_value": 0.0, "cycle_size": 12},
  },  
]

t = torch.tensor(0)


def display_create_lr_scheduler_with_warmup(configs):
  
  print("torch version:", torch.__version__)
  print("ignite master:", ignite.__version__)
  
  n = 2

  for i, config in enumerate(configs):
    warmup_duration = 4    
    optimizer = torch.optim.SGD([t, ], lr=config['init_lr'])
    lr_scheduler = config['lr_scheduler_cls'](optimizer, **config['lr_scheduler_kwargs'])
    
    scheduler = create_lr_scheduler_with_warmup(
        lr_scheduler,
        warmup_start_value=0.0,
        warmup_end_value=config['warmup_end_value'],
        warmup_duration=warmup_duration,
    )
    
    lr_values = []

    def save_lr(engine):
        lr_values.append((engine.state.iteration - 1, optimizer.param_groups[0]['lr']))

    trainer = Engine(lambda engine, batch: None)
    trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
    trainer.add_event_handler(Events.ITERATION_STARTED, save_lr)

    trainer.run(list(range(31)), max_epochs=1)
    lr_values = np.array(lr_values)
    
    if i % n == 0:
        plt.figure(figsize=(10 * n, 5))
        
    init_lr = config['init_lr'] if not isinstance(lr_scheduler, LinearCyclicalScheduler) else config['lr_scheduler_kwargs']['start_value']
        
    plt.subplot(1, n, (i % n) + 1)
    plt.title("{}\n{}\n".format(config['name'], 
                                config['description'].format(
                                  init_lr, config['warmup_end_value'], warmup_duration)))
    plt.plot(lr_values[:, 0], lr_values[:, 1], '*-', label="learning rate")
    plt.legend()
    plt.grid(which='both')
```





Discussion:

Patch:
diff --git a/ignite/contrib/handlers/param_scheduler.py b/ignite/contrib/handlers/param_scheduler.py
index 4c90f9d5..9319c613 100644
--- a/ignite/contrib/handlers/param_scheduler.py
+++ b/ignite/contrib/handlers/param_scheduler.py
@@ -423,7 +423,13 @@ class LRScheduler(ParamScheduler):
 
         step_scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
         scheduler = LRScheduler(step_scheduler)
-        trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
+
+        # In this example, we assume to have installed PyTorch>=1.1.0
+        # (with new `torch.optim.lr_scheduler` behaviour) and
+        # we attach scheduler to Events.ITERATION_COMPLETED
+        # instead of Events.ITERATION_STARTED to make sure to use
+        # the first lr value from the optimizer, otherwise it is will be skipped:
+        trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
     """
 
     def __init__(self, lr_scheduler, save_history=False, **kwds):
@@ -450,7 +456,10 @@ class LRScheduler(ParamScheduler):
     def get_param(self):
         """Method to get current optimizer's parameter value
         """
+        # Emulate context manager for pytorch>=1.4
+        self.lr_scheduler._get_lr_called_within_step = True
         lr_list = self.lr_scheduler.get_lr()
+        self.lr_scheduler._get_lr_called_within_step = False
         if len(lr_list) > 1:
             raise ValueError("Optimizer passed to lr_scheduler should have a single param group, "
                              "but currently there are {} param groups".format(len(lr_list)))
@@ -475,8 +484,8 @@ class LRScheduler(ParamScheduler):
         values = []
         scheduler = cls(save_history=False, lr_scheduler=copy_lr_scheduler)
         for i in range(num_events):
-            scheduler(engine=None)
             values.append([i, scheduler.optimizer_param_groups[0][scheduler.param_name]])
+            scheduler(engine=None)
 
         return values
 
@@ -493,7 +502,7 @@ class LRScheduler(ParamScheduler):
         for group in dummy_optimizer.param_groups:
             group.setdefault('initial_lr', group['lr'])
         kwargs = lr_scheduler.state_dict()
-        for k in ['base_lrs', '_step_count']:
+        for k in [_k for _k in kwargs.keys() if "_" == _k[0]] + ['base_lrs', 'last_epoch']:
             del kwargs[k]
         copy_lr_scheduler = lr_scheduler_cls(optimizer=dummy_optimizer, **kwargs)
         copy_lr_scheduler.load_state_dict(lr_scheduler.state_dict())
@@ -504,58 +513,82 @@ def create_lr_scheduler_with_warmup(lr_scheduler, warmup_start_value, warmup_end
                                     save_history=False,
                                     output_simulated_values=None):
     """
-    Helper method to create a LR scheduler with a linear warm-up.
+    Helper method to create a learning rate scheduler with a linear warm-up.
 
     Args:
-        lr_scheduler (ParamScheduler or subclass of `torch.optim.lr_scheduler._LRScheduler`): LR scheduler after
-            the warm-up.
-        warmup_start_value (float): LR start value of the warm-up phase.
-        warmup_end_value (float): LR end value of the warm-up phase.
+        lr_scheduler (ParamScheduler or subclass of `torch.optim.lr_scheduler._LRScheduler`): learning rate scheduler
+            after the warm-up.
+        warmup_start_value (float): learning rate start value of the warm-up phase.
+        warmup_end_value (float): learning rate end value of the warm-up phase.
         warmup_duration (int): warm-up phase duration, number of events.
         save_history (bool, optional): whether to log the parameter values to
             `engine.state.param_history`, (default=False).
-        output_simulated_values (list, optional): optional output of simulated LR values.
+        output_simulated_values (list, optional): optional output of simulated learning rate values.
             If output_simulated_values is a list of None, e.g. `[None] * 100`, after the execution it will be filled
-            by 100 simulated LR values.
+            by 100 simulated learning rate values.
 
     Returns:
-        ConcatScheduler: LR scheduler with linear warm-up.
+        ConcatScheduler: learning rate scheduler with linear warm-up.
 
+    Note:
+        If the first learning rate value provided by `lr_scheduler` is different from `warmup_end_value`, an additional
+        event is added after the warm-up phase such that the warm-up ends with `warmup_end_value` value and then
+        `lr_scheduler` provides its learning rate values as normally.
 
-    .. code-block:: python
+    Examples:
 
-        torch_lr_scheduler = ExponentialLR(optimizer=optimizer, gamma=0.98)
-        lr_values = [None] * 100
-        scheduler = create_lr_scheduler_with_warmup(torch_lr_scheduler,
-                                                    warmup_start_value=0.0, warmup_end_value=0.1, warmup_duration=10,
-                                                    output_simulated_values=lr_values)
-        lr_values = np.array(lr_values)
-        # Plot simulated values
-        plt.plot(lr_values[:, 0], lr_values[:, 1], label="learning rate")
+        .. code-block:: python
 
-        # Attach to the trainer
-        trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
+            torch_lr_scheduler = ExponentialLR(optimizer=optimizer, gamma=0.98)
+            lr_values = [None] * 100
+            scheduler = create_lr_scheduler_with_warmup(torch_lr_scheduler,
+                                                        warmup_start_value=0.0,
+                                                        warmup_end_value=0.1,
+                                                        warmup_duration=10,
+                                                        output_simulated_values=lr_values)
+            lr_values = np.array(lr_values)
+            # Plot simulated values
+            plt.plot(lr_values[:, 0], lr_values[:, 1], label="learning rate")
+
+            # Attach to the trainer
+            trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
 
     """
     if not isinstance(lr_scheduler, (ParamScheduler, _LRScheduler)):
         raise TypeError("Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler._LRScheduler or "
                         "ParamScheduler, but given {}".format(type(lr_scheduler)))
 
+    if not (isinstance(warmup_duration, numbers.Integral) and warmup_duration > 1):
+        raise ValueError("Argument warmup_duration should be at least 2 events, but given {}"
+                         .format(warmup_duration))
+
+    milestones_values = [(0, warmup_start_value), (warmup_duration - 1, warmup_end_value)]
+
     duration_extension = 0
     if isinstance(lr_scheduler, _LRScheduler):
+        init_lrs = [g['lr'] for g in lr_scheduler.optimizer.param_groups]
+        if len(init_lrs) < 1:
+            raise RuntimeError("Number of parameter groups of input `lr_scheduler.optimizer` is less than one.")
+
+        if init_lrs[0] != warmup_end_value:
+            milestones_values.append((warmup_duration, init_lrs[0]))
+
         lr_scheduler = LRScheduler(lr_scheduler)
-        duration_extension = 1
+    else:
+        init_lr = lr_scheduler.get_param()
+        if init_lr == warmup_end_value:
+            if warmup_duration > 2:
+                d = (warmup_end_value - warmup_start_value) / (warmup_duration - 1)
+                milestones_values[-1] = (warmup_duration - 2, warmup_end_value - d)
+            else:
+                milestones_values.pop(-1)
 
     dummy_optimizer = {}
-    warmup_scheduler = LinearCyclicalScheduler(dummy_optimizer, param_name="lr",
-                                               start_value=warmup_start_value,
-                                               end_value=warmup_end_value,
-                                               cycle_size=warmup_duration * 2)
-
+    warmup_scheduler = PiecewiseLinear(dummy_optimizer, param_name="lr", milestones_values=milestones_values)
     warmup_scheduler.optimizer_param_groups = lr_scheduler.optimizer_param_groups
 
     schedulers = [warmup_scheduler, lr_scheduler]
-    durations = [warmup_duration + duration_extension, ]
+    durations = [milestones_values[-1][0] + 1, ]
     combined_scheduler = ConcatScheduler(schedulers, durations=durations,
                                          save_history=save_history)
     if output_simulated_values is not None:


Test patch:
diff --git a/tests/ignite/contrib/handlers/test_param_scheduler.py b/tests/ignite/contrib/handlers/test_param_scheduler.py
index 7009d311..bfbcf3f8 100644
--- a/tests/ignite/contrib/handlers/test_param_scheduler.py
+++ b/tests/ignite/contrib/handlers/test_param_scheduler.py
@@ -363,18 +363,18 @@ def test_lr_scheduler():
 
         trainer = Engine(dummy_update)
 
-        @trainer.on(Events.ITERATION_COMPLETED)
-        def torch_lr_scheduler_step(engine):
-            torch_lr_scheduler2.step()
-
-        @trainer.on(Events.ITERATION_COMPLETED)
+        @trainer.on(Events.ITERATION_STARTED)
         def save_lr(engine):
             lrs.append(optimizer1.param_groups[0]['lr'])
 
-        @trainer.on(Events.ITERATION_COMPLETED)
+        @trainer.on(Events.ITERATION_STARTED)
         def save_true_lr(engine):
             lrs_true.append(optimizer2.param_groups[0]['lr'])
 
+        @trainer.on(Events.ITERATION_COMPLETED)
+        def torch_lr_scheduler_step(engine):
+            torch_lr_scheduler2.step()
+
         trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
 
         data = [0] * 10
@@ -470,9 +470,11 @@ def test_simulate_values():
     def _test(scheduler_cls, **scheduler_kwargs):
 
         optimizer = None
+        event = Events.ITERATION_STARTED
         if scheduler_cls == LRScheduler:
             scheduler_kwargs['optimizer'] = scheduler_kwargs['lr_scheduler'].optimizer
             optimizer = scheduler_kwargs['optimizer']
+            event = Events.ITERATION_COMPLETED
         elif scheduler_cls == ConcatScheduler:
             optimizer = scheduler_kwargs['optimizer']
             del scheduler_kwargs['optimizer']
@@ -493,8 +495,8 @@ def test_simulate_values():
             lrs.append(optimizer.param_groups[0]['lr'])
 
         trainer = Engine(lambda engine, batch: None)
-        trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
-        trainer.add_event_handler(Events.ITERATION_COMPLETED, save_lr)
+        trainer.add_event_handler(event, scheduler)
+        trainer.add_event_handler(Events.ITERATION_STARTED, save_lr)
         trainer.run(data, max_epochs=max_epochs)
 
         assert lrs == pytest.approx([v for i, v in simulated_values])
@@ -545,19 +547,35 @@ def test_simulate_values():
 
 
 def test_create_lr_scheduler_with_warmup():
-    with pytest.raises(TypeError):
+
+    with pytest.raises(TypeError, match=r"Argument lr_scheduler should be a subclass of"):
         create_lr_scheduler_with_warmup(12, warmup_start_value=0.0, warmup_end_value=0.1, warmup_duration=10)
 
-    def _test(lr_scheduler, optimizer, warmup_end_value_to_check=None):
+    t1 = torch.zeros([1], requires_grad=True)
+    # A) opt lr != warmup_end_value
+    optimizer = torch.optim.SGD([t1], lr=0.2)
+    torch_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.98)
+
+    with pytest.raises(ValueError, match=r"Argument warmup_duration should be at least 2 events"):
+        create_lr_scheduler_with_warmup(torch_lr_scheduler, warmup_start_value=0.0,
+                                        warmup_end_value=0.1, warmup_duration=1)
+
+    with pytest.raises(ValueError, match=r"Argument warmup_duration should be at least 2 events"):
+        create_lr_scheduler_with_warmup(torch_lr_scheduler, warmup_start_value=0.0,
+                                        warmup_end_value=0.1, warmup_duration="abc")
+
+    with pytest.raises(TypeError, match=r"Argument output_simulated_values should be a list of None"):
+        simulated_values = ()
+        create_lr_scheduler_with_warmup(torch_lr_scheduler, warmup_start_value=0.0, warmup_end_value=0.1,
+                                        warmup_duration=10, output_simulated_values=simulated_values)
+
+    def _test(lr_scheduler, optimizer, warmup_start_value, warmup_end_value, warmup_duration, warmup_end_next_value):
         num_iterations = 10
         max_epochs = 20
 
-        warmup_duration = 10
-        warmup_end_value = 0.1
-
         simulated_values = [None] * (num_iterations * max_epochs)
         scheduler = create_lr_scheduler_with_warmup(lr_scheduler,
-                                                    warmup_start_value=0.0,
+                                                    warmup_start_value=warmup_start_value,
                                                     warmup_end_value=warmup_end_value,
                                                     warmup_duration=warmup_duration,
                                                     output_simulated_values=simulated_values)
@@ -565,31 +583,61 @@ def test_create_lr_scheduler_with_warmup():
         lrs = []
         trainer = Engine(lambda engine, batch: None)
 
-        @trainer.on(Events.ITERATION_COMPLETED)
+        trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
+        @trainer.on(Events.ITERATION_STARTED)
         def save_lr(engine):
             lrs.append(optimizer.param_groups[0]['lr'])
 
-        trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
-
         data = [0] * num_iterations
         trainer.run(data, max_epochs=max_epochs)
 
         assert lrs == pytest.approx([v for i, v in simulated_values])
 
-        if warmup_end_value_to_check is None:
-            warmup_end_value_to_check = warmup_end_value
-        assert lrs[warmup_duration] == warmup_end_value_to_check
+        assert lrs[0] == pytest.approx(warmup_start_value), \
+            "lrs={}".format(lrs[:warmup_duration + num_iterations])
+        assert lrs[warmup_duration - 1] == pytest.approx(warmup_end_value), \
+            "lrs={}".format(lrs[:warmup_duration + num_iterations])
+        assert lrs[warmup_duration] == pytest.approx(warmup_end_next_value), \
+            "lrs={}".format(lrs[:warmup_duration + num_iterations])
 
     t1 = torch.zeros([1], requires_grad=True)
+    # A) opt lr != warmup_end_value
+    optimizer = torch.optim.SGD([t1], lr=0.2)
+    torch_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.98)
+    _test(torch_lr_scheduler, optimizer, 0.01, 0.05, 10, 0.2)
     optimizer = torch.optim.SGD([t1], lr=0.2)
     torch_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.98)
-    _test(torch_lr_scheduler, optimizer)
+    _test(torch_lr_scheduler, optimizer, 0.01, 0.05, 2, 0.2)
 
-    t1 = torch.zeros([1], requires_grad=True)
+    # B) opt lr == warmup_end_value
+    optimizer = torch.optim.SGD([t1], lr=0.2)
+    torch_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.98)
+    _test(torch_lr_scheduler, optimizer, 0.01, 0.2, 10, 0.2 * 0.98)
     optimizer = torch.optim.SGD([t1], lr=0.2)
+    torch_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.98)
+    _test(torch_lr_scheduler, optimizer, 0.01, 0.2, 2, 0.2 * 0.98)
+
+    # C) lr_scheduler start_value != warmup_end_value
+    t1 = torch.zeros([1], requires_grad=True)
+    optimizer = torch.optim.SGD([t1], lr=0.0)
+    lr_scheduler = LinearCyclicalScheduler(optimizer=optimizer, param_name='lr',
+                                           start_value=0.8, end_value=0.0, cycle_size=10)
+    _test(lr_scheduler, optimizer, 0.01, 0.05, 10, 0.8)
+    optimizer = torch.optim.SGD([t1], lr=0.0)
     lr_scheduler = LinearCyclicalScheduler(optimizer=optimizer, param_name='lr',
-                                           start_value=1.0, end_value=0.0, cycle_size=10)
-    _test(lr_scheduler, optimizer, 1.0)
+                                           start_value=0.8, end_value=0.0, cycle_size=10)
+    _test(lr_scheduler, optimizer, 0.01, 0.05, 2, 0.8)
+
+    # D) lr_scheduler start_value == warmup_end_value
+    t1 = torch.zeros([1], requires_grad=True)
+    optimizer = torch.optim.SGD([t1], lr=0.0)
+    lr_scheduler = LinearCyclicalScheduler(optimizer=optimizer, param_name='lr',
+                                           start_value=0.8, end_value=0.0, cycle_size=10)
+    _test(lr_scheduler, optimizer, 0.01, 0.8, 10, 0.8 - (0.8 / 5.0))
+    optimizer = torch.optim.SGD([t1], lr=0.0)
+    lr_scheduler = LinearCyclicalScheduler(optimizer=optimizer, param_name='lr',
+                                           start_value=0.8, end_value=0.0, cycle_size=10)
+    _test(lr_scheduler, optimizer, 0.01, 0.8, 2, 0.8 - (0.8 / 5.0))
 
 
 def test_create_lr_scheduler_with_warmup_on_combined_scheduler():
@@ -715,19 +763,24 @@ def test_param_group_scheduler():
 def test_create_lr_scheduler_with_warmup_with_real_model(dummy_model_factory):
 
     model = dummy_model_factory(with_grads=False, with_frozen_layer=False)
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    init_lr = 0.01
+    optimizer = torch.optim.SGD(model.parameters(), lr=init_lr)
     scaled_lr = 0.02
+    warmup_duration = 5
+    step_size = 2
+    gamma = 0.97
 
     output_simulated_values = [None] * 50
 
     create_lr_scheduler_with_warmup(
-        torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.97),
+        torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma),
         warmup_start_value=0.0,
         warmup_end_value=scaled_lr,
-        warmup_duration=5,
+        warmup_duration=warmup_duration,
         output_simulated_values=output_simulated_values
     )
-
     assert output_simulated_values[0] == [0, 0.0]
-    assert output_simulated_values[5] == [5, scaled_lr]
-    assert output_simulated_values[6] == [5 + 1, 0.01]
+    assert output_simulated_values[warmup_duration - 1] == [warmup_duration - 1, scaled_lr]
+    assert output_simulated_values[warmup_duration] == [warmup_duration, init_lr]
+    v = [warmup_duration + step_size, init_lr * gamma]
+    assert output_simulated_values[warmup_duration + step_size] == v