Discriminative learning rates in NeMo? [Question]
State: closed | Pull request: no

Is it possible to use discriminative learning rates with NeMo, i.e. a different lr for different layer groups like the encoder and decoder? This is a technique that can be very helpful for transfer learning and I would like to experiment with it if possible. I tried passing a list of LRs to the optimizer config but it seems to be hardcoded for a single LR.

![image](https://user-images.githubusercontent.com/47190785/99979815-87574880-2d75-11eb-955a-9eb6655b68aa.png)

**Additional context**

I'm using quartznet with nemo-toolkit[all]==1.0.0b2


It's not supported out of the box, but you can easily add support for it by subclassing EncDecCTCModel and making the required modifications. 

So I finally got to work on this and I may need some help as I don't understand all the stuff under the hood. Here's what my subclass of `EncDecCTCModel` looks like. Code adopted from [here](https://github.com/PyTorchLightning/pytorch-lightning/issues/2005)

```
from omegaconf import DictConfig
from pytorch_lightning import Trainer

class EncDecSubclass(nemo_asr.models.EncDecCTCModel):
    def __init__(self, cfg: DictConfig, trainer:pl.Trainer=None):
        super(EncDecSubclass, self).__init__(cfg, trainer)

    def configure_optimizers(self):
        params = list(self.named_parameters())
        def is_backbone(n): return 'encoder' in n
        
        # Is this the appropriate way to access the LR from inside configure optimizers? 
        base_lr = self.cfg.optim.lr
        grouped_parameters = [
            {'params': [p for n, p in params if is_backbone(n)], 'lr': base_lr/100},
            {'params': [p for n, p in params if not is_backbone(n)], 'lr': base_lr},
        ]

        optimizer = torch.optim.AdamW(
            grouped_parameters, lr=1e-3, weight_decay=0
        )
        #Is this right? I stored it here to avoid an error I was getting
        self._optimizer = optimizer
        return optimizer
```

Then I instantiate the model using the quartznet15x5 config and restore a copy of quartznet15x5 weights.
```
empty_model = EncDecSubclass(cfg=DictConfig(params['model']))
my__model = empty_model.restore_from(PATH_SAVED_MODELS/'quartznet_pretrained.nemo')
my_model.setup_training_data(train_data_config=params["model"]["train_ds"])
my_model.setup_validation_data(val_data_config=params["model"]["validation_ds"])
freeze(my_model) # code to turn off autograd on encoder params
trainer = pl.Trainer(gpus=1, max_steps=2000)
my__model.set_trainer(trainer)
trainer.fit(my__model)
```

I then unfreeze, set the LR with `my_model.setup_optimization` and continue training. Everything seems to be working. My questions are :

1. Does anyone see anything that jumps out to them as wrong here? 
2. How can I verify it is actually using the discriminative learning rates? My thought is to try setting the lr for the param groups for something really high (like 1000) for the encoder, and a normal level for the decoder, and make sure it trains well when frozen and then diverges when unfrozen? **Edit:** I tried this and it worked! It doesn't cause the lr_finder to diverge though so 
the discriminative learning rates aren't propogating to the lr finder. Any ideas here? 

Hope this code helps someone! Thank you. 


Hey @rbracco, for (1), yes - this looks good. This is a perfectly valid way to use a modified optimization workflow. 

While the setup_optimization() code does a lot of stuff inside to properly calculate schedule and instantiate the correct Optimizer and scheduler from yaml, the configure_optimizer call you have overriden here simply calls setup_optimization and then returns `self._optimizer`.

So this approach is perfectly fine, but you will lose the ability to instantiate or change optimizer from config since this bypasses setup_optimization. 

We don't do any magic with the optimizers themselves - they behave just as any other pytorch optimizer behaves, so you will get discriminative LR here. 

For (2), logging will actually display the LR of the first param group - aka the backbone lr (lr / 100). So if you find that your lr starts off from (lr / 100) on your logs, then it is working properly.