FunASR/funasr/schedulers/lambdalr_cus.py
zhifu gao 4a7a984a5f
Dev gzf (#1465)
* qwenaudio qwenaudiochat

* qwenaudio qwenaudiochat

* whisper

* whisper

* llm

* llm

* llm

* llm

* llm

* llm

* llm

* llm

* export onnx

* export onnx

* export onnx

* dingding

* dingding

* llm
2024-03-11 17:56:30 +08:00

17 lines
546 B
Python

import torch
from torch.optim.lr_scheduler import _LRScheduler
class CustomLambdaLR(_LRScheduler):
def __init__(self, optimizer, warmup_steps, last_epoch=-1):
self.warmup_steps = warmup_steps
super().__init__(optimizer, last_epoch)
def get_lr(self):
if self.last_epoch < self.warmup_steps:
return [
base_lr * min(self.last_epoch / self.warmup_steps, 1)
for base_lr in self.base_lrs
]
else:
return [base_lr for base_lr in self.base_lrs]