Make LRScheduler attachable to Events.ITERATION_STARTED

[vfdev-5]

🚀 Feature

Currently, the correct way to use LRScheduler wrapper for pytorch >= 1.1.0 is the following:

from torch.optim.lr_scheduler import StepLR

torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
scheduler = LRScheduler(torch_lr_scheduler)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])

# 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)

trainer..run([0] * 8, max_epochs=1)

0.1
0.1
0.1
0.010
0.010
0.010
0.001
0.001

• Fixed docs example for LRScheduler #2463

however, other schedulers should be used as following (link)

milestones_values = [(1, 1.0), (3, 0.8), (5, 0.2)]
scheduler = PiecewiseLinear(
    optimizer, "lr", milestones_values=milestones_values)

trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])

trainer.run([0] * 6, max_epochs=1)

The idea is to improve LRScheduler such that we could attach it to Events.ITERATION_STARTED and have a coherent API. It will be a BC-breaking change, but for good.

So, desired example using LRScheduler should be:

from torch.optim.lr_scheduler import StepLR

torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
scheduler = LRScheduler(torch_lr_scheduler)

trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)

@trainer.on(Events.ITERATION_COMPLETED)
def print_lr():
    print(optimizer.param_groups[0]["lr"])

trainer.run([0] * 8, max_epochs=1)

Currently, this gives a wrong behaviour as the first 0.1 wasn't consumed by the training step.

0.1
0.1
0.010
0.010
0.010
0.001
0.001

The idea could be to retain the first value and reapply it once and then keep everything as it is now.

[sdesrozis]

Maybe an extra argument could be introduced to keep the current (but not expected) behaviour ?

[avats-dev]

I would like to work on implementing this. Can you provide me some reference for Events.ITERATION_STARTED and Events.ITERATION_COMPLETED ?

[vfdev-5]

@avats-dev checkout https://pytorch-ignite.ai/concepts/02-events-and-handlers/

[yuta0821]

Hi, I would like to contribute about this issue.
I have read through the source code for the relevant parts.
If don't use the current way to use LRScheduler wrapper for pytorch >= 1.1.0, I think that we can maintain the original output by replacing the code as follows

• current code
  

  ignite/ignite/handlers/param_scheduler.py

  Lines 858 to 859 in 493d86d

  	self.lr_scheduler.last_epoch += 1 # type: ignore[attr-defined]
  	super(LRScheduler, self).__call__(engine, name)

• suggested code

super(LRScheduler, self).__call__(engine, name)
self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]

• output
  • current example

  from torch.optim.lr_scheduler import StepLR
  torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1) 
  scheduler = LRScheduler(torch_lr_scheduler)
  
  @trainer.on(Events.ITERATION_COMPLETED)
  def print_lr():
      print(optimizer.param_groups[0]["lr"])
  
  trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
  trainer.run([0] * 8, max_epochs=1)

  0.1
  0.1
  0.1
  0.1
  0.010
  0.010 
  0.010
  0.001
  

  • desired example using LRScheduler

  from torch.optim.lr_scheduler import StepLR
  torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1) 
  scheduler = LRScheduler(torch_lr_scheduler)
  
  trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
  
  @trainer.on(Events.ITERATION_COMPLETED)
  def print_lr():
      print(optimizer.param_groups[0]["lr"])
  
  trainer.run([0] * 8, max_epochs=1)

  0.1
  0.1
  0.1
  0.010
  0.010 
  0.010
  0.001
  0.001
  

If it is necessary to make it compatible with the current example, we probably need to add other functions.
I'm sorry if my opinion is misguided.

[sdesrozis]

It sounds good !

Maybe you could just add an extra argument use_legacy or whatever and a suitable depreciation message if we would have the current behaviour for BC.

Well done though !

[yuta0821]

@sdesrozis
Thanks a lot !

OK, can I submit a PR with suggested code and extra argument use_legacy: bool with a warning message?

[yuta0821]

Based on vfdev-5's idea, I have come up with a better code using update_last_epoch, which I propose below.

• code

  class LRScheduler(ParamScheduler):
    def __init__(self, lr_scheduler: _LRScheduler, save_history: bool = False, update_last_epoch: bool = False):
  
        self.lr_scheduler = lr_scheduler
        super(LRScheduler, self).__init__(
            optimizer=self.lr_scheduler.optimizer,  # type: ignore[attr-defined]
            param_name="lr",
            save_history=save_history,
        )
        self.update_last_epoch = update_last_epoch
        self._state_attrs += ["lr_scheduler"]
  
    def __call__(self, engine: Optional[Engine], name: Optional[str] = None) -> None:
        if self.update_last_epoch:
            self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]
        else:
            self.update_last_epoch = True
        super(LRScheduler, self).__call__(engine, name)
  
    def get_param(self) -> Union[float, List[float]]:
        ...
  
    @classmethod
    def simulate_values(  # type: ignore[override]
        cls, num_events: int, lr_scheduler: _LRScheduler, **kwargs: Any
    ) -> List[List[int]]:
        with tempfile.TemporaryDirectory() as tmpdirname:
            cache_filepath = Path(tmpdirname) / "ignite_lr_scheduler_cache.pt"
            obj = {
                "lr_scheduler": lr_scheduler.state_dict(),
                "optimizer": lr_scheduler.optimizer.state_dict(),  # type: ignore[attr-defined]
            }
            torch.save(obj, cache_filepath.as_posix())
  
            values = []
            scheduler = cls(save_history=False, lr_scheduler=lr_scheduler, update_last_epoch=True, **kwargs)
            for i in range(num_events):
                params = [p[scheduler.param_name] for p in scheduler.optimizer_param_groups]
                values.append([i] + params)
                scheduler(engine=None)
  
            obj = torch.load(cache_filepath.as_posix())
            lr_scheduler.load_state_dict(obj["lr_scheduler"])
            lr_scheduler.optimizer.load_state_dict(obj["optimizer"])  # type: ignore[attr-defined]
  
            return values

• output
  • current example

  from torch.optim.lr_scheduler import StepLR
  torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1) 
  scheduler = LRScheduler(torch_lr_scheduler, update_last_epoch=True)
  
  @trainer.on(Events.ITERATION_COMPLETED)
  def print_lr():
      print(optimizer.param_groups[0]["lr"])
  
  trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
  trainer.run([0] * 8, max_epochs=1)

  0.1
  0.1
  0.1
  0.010
  0.010 
  0.010
  0.001
  0.001
  

  • desired example using LRScheduler

  from torch.optim.lr_scheduler import StepLR
  torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1) 
  scheduler = LRScheduler(torch_lr_scheduler)
  
  trainer.add_event_handler(Events.ITERATION_STARTED, scheduler)
  
  @trainer.on(Events.ITERATION_COMPLETED)
  def print_lr():
      print(optimizer.param_groups[0]["lr"])
  
  trainer.run([0] * 8, max_epochs=1)

  0.1
  0.1
  0.1
  0.010
  0.010 
  0.010
  0.001
  0.001
  

  • simulate values example using LRScheduler

  from torch.optim.lr_scheduler import StepLR
  torch_lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1) 
  LRScheduler.simulate_values(10, torch_lr_scheduler)

  [[0, 0.0001],
   [1, 0.0001],
   [2, 0.0001],
   [3, 1e-05],
   [4, 1e-05],
   [5, 1e-05],
   [6, 1.0000000000000002e-06],
   [7, 1.0000000000000002e-06],
   [8, 1.0000000000000002e-06],
   [9, 1.0000000000000002e-07]]
  

[vfdev-5]

@yuta0821 if I understand correctly, you introduce update_last_epoch flag and you use it inside __call__ method.
If we could have a better name, it would be better from user experience. Secondly, I think if we need to call once more self.lr_scheduler.last_epoch += 1 for the previous behaviour and just do

super(LRScheduler, self).__call__(engine, name)
self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]

for the new behaviour, let's call self.lr_scheduler.last_epoch += 1 inside the constructor according the flag.
What do you think ?

[yuta0821]

@vfdev-5
Thank you for your advice.
It sounds good !
If put your idea into code, is it corret to express it as follows ?

• code

  class LRScheduler(ParamScheduler):
    def __init__(self, lr_scheduler: _LRScheduler, save_history: bool = False, use_legacy: bool = False):
        """
        Args:
            use_legacy: the flag for the previous behaviour to 
                decide whether to call once more 'self.lr_scheduler.last_epoch += 1'
        """
  
        self.lr_scheduler = lr_scheduler
        super(LRScheduler, self).__init__(
            optimizer=self.lr_scheduler.optimizer,  # type: ignore[attr-defined]
            param_name="lr",
            save_history=save_history,
        )
        self.use_legacy = use_legacy 
        if use_legacy:
            # add some warning message about the scheduler needing to be attached to Events.ITERATION_COMPLETED
            warnings.warn("we attach scheduler to Events.ITERATION_COMPLETED ...")
        self._state_attrs += ["lr_scheduler"]
  
    def __call__(self, engine: Optional[Engine], name: Optional[str] = None) -> None:
        # call once more 'self.lr_scheduler.last_epoch += 1' for the previous behaviour
        if self.use_legacy: 
            self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]
            self.use_legacy = False
        # just do for the new behaviour
        super(LRScheduler, self).__call__(engine, name)
        self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]
  
    def get_param(self) -> Union[float, List[float]]:
        ...
  
    @classmethod
    def simulate_values(  # type: ignore[override]
        cls, num_events: int, lr_scheduler: _LRScheduler, **kwargs: Any
    ) -> List[List[int]]:
        with tempfile.TemporaryDirectory() as tmpdirname:
            cache_filepath = Path(tmpdirname) / "ignite_lr_scheduler_cache.pt"
            obj = {
                "lr_scheduler": lr_scheduler.state_dict(),
                "optimizer": lr_scheduler.optimizer.state_dict(),  # type: ignore[attr-defined]
            }
            torch.save(obj, cache_filepath.as_posix())
  
            values = []
            scheduler = cls(save_history=False, lr_scheduler=lr_scheduler, **kwargs)
            for i in range(num_events):
                scheduler(engine=None) # adjust the timing of the scheduler call according to changes in the __call__ method
                params = [p[scheduler.param_name] for p in scheduler.optimizer_param_groups]
                values.append([i] + params)
  
            obj = torch.load(cache_filepath.as_posix())
            lr_scheduler.load_state_dict(obj["lr_scheduler"])
            lr_scheduler.optimizer.load_state_dict(obj["optimizer"])  # type: ignore[attr-defined]
  
            return values

[vfdev-5]

@yuta0821 thanks for the code. The idea I exposed previously was to avoid if statement in the __call__ and do something like:

class LRScheduler(ParamScheduler):
  def __init__(self, lr_scheduler: _LRScheduler, save_history: bool = False, use_legacy: bool = False):
      """
      Args:
          use_legacy: the flag for the previous behaviour to 
              decide whether to call once more 'self.lr_scheduler.last_epoch += 1'
      """

      self.lr_scheduler = lr_scheduler
      super(LRScheduler, self).__init__(
          optimizer=self.lr_scheduler.optimizer,  # type: ignore[attr-defined]
          param_name="lr",
          save_history=save_history,
      )
      if use_legacy:
          # add some warning message about the scheduler needing to be attached to Events.ITERATION_COMPLETED
          warnings.warn("we attach scheduler to Events.ITERATION_COMPLETED ...")
          self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]

      self._state_attrs += ["lr_scheduler"]

  def __call__(self, engine: Optional[Engine], name: Optional[str] = None) -> None:
      # just do for the new behaviour
      super(LRScheduler, self).__call__(engine, name)
      self.lr_scheduler.last_epoch += 1  # type: ignore[attr-defined]

Can you please check if this works ?

[yuta0821]

@vfdev-5
OK, I understand.
I check if your code works and it works well.

[yuta0821]

@vfdev-5
Can I submit PR about this improvement ?