Bug in CosineAnnealingWarmRestarts

[heilaw]

🐛 Bug

CosineAnnealingWarmRestarts gives AttributeError if we initialize it with last_epoch not equal to -1, which usually happens when one wants to resume training a network.

To Reproduce

Steps to reproduce the behavior:

1. Initialize the scheduler and set the last_epoch to where we left off

scheduler = CosineAnnealingWarmRestarts(
      optimizer, restart_iter,
      T_mult=restart_mult,
      last_epoch=_last_iter # where we left off, _last_epoch should be greater than -1
)

2. The code gives this error:

  File "/foo/torch/optim/lr_scheduler.py", line 965, in __init__
    super(CosineAnnealingWarmRestarts, self).__init__(optimizer, last_epoch, verbose)
  File "/foo/torch/optim/lr_scheduler.py", line 79, in __init__
    self.step()
  File "/foo/torch/optim/lr_scheduler.py", line 1008, in step
    self.T_cur = self.T_cur + 1
AttributeError: 'CosineAnnealingWarmRestarts' object has no attribute 'T_cur'

Expected behavior

No attribute error.

Environment

Please copy and paste the output from our
environment collection script
(or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

• PyTorch Version (e.g., 1.0): 1.7.0
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): conda
• Build command you used (if compiling from source):
• Python version: 3.7
• CUDA/cuDNN version: 11/6.0.21
• GPU models and configuration: RTX 3090
• Any other relevant information:

Additional context

After some investigation, I believe I found the code that causes this bug. When we initialize CosineAnnealingWarmRestarts, it will call its parent (_LRScheduler) constructor (line 965 in torch/optim/lr_scheduler.py). The constructor of _LRScheduler will call the step function (line 79 in torch/optim/lr_scheduler.py) of CosineAnnealingWarmRestarts before finishing the initialization. In the step function, because epoch is none and last_epoch is greater than 0, line 1008 will be executed which involves self.T_cur. But at this point self.T_cur has not been set yet as it is only set after the initialization is done (line 967).

cc @vincentqb

[VitalyFedyunin]

Triage review because we probably need new 'scheduler' label

[VitalyFedyunin]

Optimizer label just fine for it

[mruberry]

Thank you for reporting this issue, @heilaw. I think we would accept a PR fixing this.

[ziyuang]

Not merged yet?

[hugefrog]

bug is still exsit in pytorch 1.10