Repository: pytorch
Issue: Multiple Learning Rate Scheduler for Specific Parameters Groups
Labels: module: optimizer,triaged,needs research,has workaround,module: LrScheduler

### 🚀 The feature, motivation and pitch

The concept of learning rate schedulers is currently to affect the learning rate for all parameter groups. However there are use cases where it is needed to affect only specific parameters groups, have different parameters for the scheduler between the parameters groups, or even have different types of schedulers for each parameters group.

To give an example:
1. I want to have a warm-up phase (e.g. LinearLR from 0.1 to 1.0 over 10 epochs) but only for the network backend or head. So, only for a specific parameters group.
2. I want to have a warm-up phase for my model but with different lengths (total_iters=10 for the backend and total_iters=30 for the head or similar).
3. I want to have a warm-up phase but with different lengths and styles (e.g. linear LR for the backend for the first 10 epochs and exponential increasing LR for the head for the first 30 epochs. 

Currently, it is not possible to implement it by using LRSchedulers because they affect all parameter groups which makes it impossible to also use multiple LRSchedulers for different parameter groups or even different scheduler parameters.

### Alternatives

I see one possible solution which covers all three problems described above. We can add an additional parameter to the LRScheduler base class that specifies the parameter groups that should be affected. This would allow different parameters between the parameter groups by initializing multiple schedulers of the same type, schedulers affecting only specific parameter groups, and also multiple schedulers of different types and parameter groups.

The problem is that parameter groups are currently not specified by some parameter (e.g. name). However, properties, like name, are addable. I would suggest specifying the parameter groups that should be affected by a list of names which is optional. If no list is passed, the scheduler would affect all parameter groups as usual (backward compatibility).

Here is an example:
```python
optim = Adam([
        dict(name='backend', params=backend_model.parameters(), ...),
        dict(name='head', params=head_model.parameters(), ...),
    ], lr=0.001, ...
)

scheduler_backend = LinearLR(optim, total_iters=10, params_groups=['backend'])
scheduler_backend = LinearLR(optim, total_iters=30, params_groups=['head'])
```

I currently have a use case for that problem (see additional context). So, I am interested in a solution as soon as possible 😄 . I also can implement a solution but would prefer to discuss it before with the community.

Additional solutions:
- Updating all LRSchedulers to allow different parameters for each parameters group (many code changes, solving not all problems described above).
- Allowing scheduler parameters within the parameters group dict (e.g. `dict(params=model.parameters(), total_iters=30)`).
- More?

### Additional context

I currently work on a model that should forecast energy time-series by using different modalities (historical energy, weather, and calendar data). Separated the model into an encoding part, modal fusion, and forecasting part. If I train all layers at once, some of my layers have "frozen" neurons. So, they do not change much during training which makes no sense (it is not dying ReLU). I solved the problem by training the model step-wise (and also got much better forecasts. So, I first train the first few layers, then add more, and finally train the whole model. This is similar to having way smaller learning rates in the last layers which I also tried to implement using torch.optim LRScheduler but it is not possible without implementing a specific learning rate scheduler for my specific use case. Here is an example of what kind of learning rate scheduling I want:

![image](https://github.com/pytorch/pytorch/assets/12863612/8a0b2401-93dc-46e4-b588-f4d0d32541ac)

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar