mmrazor
mmrazor copied to clipboard
[Feature] Support distillation loss weight scheduler
Modification
- Support Cos / Linear / MultiStep loss weight scheduler. Notes: Linear one can be used for distillation loss weight warmup and MultiStep can be used for early stop.
- Add loss weight scheduler manager to record the current / base (for cos schedule) loss weight.
- Add loss weight scheduler hook.
- Before run: Build schedulers. And set the milestones according to the begin / end iter of each scheduler.
- Before train epoch: Set the current loss weight attribute of the loss weight scheduler manager according to those schedulers whose
by_epoch
attribute is True. And set the base loss weight of the loss weight scheduler manager to the current loss weight if the current iter is in the milestones. For example, suppose we want to use Linear in the first 5 epochs for warmup and then use Cosine. Before the 6th epoch, we should set the base loss weight to the current one for the calculation in the Cosine scheduler. - Before train iter: Set the current loss weight attribute of the loss weight scheduler manager according to those schedulers whose
by_epoch
attribute is False. And set the base loss weight of the loss weight scheduler manager to the current loss weight if the current iter is in the milestones.
- Add the corresponding pytest.
TODO
- Docstring