Repository: pytorch
Issue: CosineAnnealingWarmRestarts with initial warm up and weight decay applied on consecutive cycles without warm up
Labels: feature,module: optimizer,triaged,needs research

### 🚀 The feature, motivation and pitch

So far there is no optimizer in torch.optim or open-source implementations that support CosineAnnealinWarmRestarts with initial linear warmup up to n specified steps followed by CosineAnnealinWarmRestarts without any further warmup with each restart, which gets followed by weight decay.

### Initial_Warmup_Cosine_Annealing_With_Weight_Decay
![initial_warm_up_cosine_weight_ecay](https://user-images.githubusercontent.com/56112578/175832056-e66ff079-cdb0-47f0-bfbb-9b1c4846ae60.png)

### Initial_Warmup_Without_Weight_Decay
![initial_warmup_withou_weight_decay](https://user-images.githubusercontent.com/56112578/175832063-31b7b5b7-f447-4e97-a7eb-2de483adbca6.png)

### No_Initial_Warmup_With_Weight_Decay
![no_initial_warm_up_with Weight_decay](https://user-images.githubusercontent.com/56112578/175832075-c177186d-ae33-40a5-8837-74b26c8d4cb2.png)


### Alternatives

Alternatives involve the ChainedScheduler paradigm which is most suitable for mutex schedulers. In order to achieve this feature, I followed the high-level design pattern of ChainedScheduler and further simplified the implementation so that it is clean, easier to understand, and above all alleviated the need to initialize multiple schedules ahead of time, before passing them to the chained scheduler which can be extra painful in this use case i.e. state staring between the two, required to pass the same optimizer to both schedules that can be potential pitfalls among others if not done correctly. 

Final Thoughts:
My approach provides a much simple, modular solution that provides the flexibility to tweak individual components to generate much more interesting learning rates that can come in handy while experimenting and quickly iterating through different variations of the same lr_scheduler.
As an example of how ChainedScheduler would look like for the case where two different schedulers need to have synchronized and updated states that are to be used by the next scheduler:
` ``
       >>> warmup_scheduler = WarmUpScheduler(
        >>>    optimizer,
        >>>    T_0=20,
        >>>    T_mul=1,
        >>>    eta_min=1e-5,
        >>>    warmup_steps=10,
        >>>    max_lr=1.0,
        >>>    gamma=1.0,
        >>> )
        >>> cosine_warm_restarts_decay = CosineAnealingWarmRestartsWeightDecay(
        >>>    optimizer,
        >>>    T_0=20,
        >>>    T_mul=1,
        >>>    eta_min=1e-5,
        >>>    warmup_steps=0,
        >>>    max_lr=1.0,
        >>>    gamma=0.9,
        >>> )
        >>> model = AlexNet(num_classes=2)
        >>> optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-1)
        >>> scheduler = ChainedScheduler([warmup_scheduler, cosine_warm_restarts_decay], warmup_steps=10)
        >>> for epoch in range(100):
        >>>     train(...)
        >>>     validate(...)
        >>>     optimizer.step()
        >>>     scheduler.step()
```

### Additional context

CosineAnealingWarmRestarts with linear warm-up during the initial steps (strictly for the number of steps specified during early stages) followed by weight decay is an extremely powerful concept and a life saver at times when training deep models as the learning rate directly affects the time required for a training run, compute resources, model convergence to either local optimum or a global one. (just to name a few)

cc @vincentqb @jbschlosser @albanD