|
|
"""
|
|
|
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
|
|
|
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
|
|
|
"""
|
|
|
|
|
|
from torch.optim.lr_scheduler import LRScheduler
|
|
|
|
|
|
from ..core import register
|
|
|
|
|
|
|
|
|
class Warmup(object):
|
|
|
def __init__(
|
|
|
self, lr_scheduler: LRScheduler, warmup_duration: int, last_step: int = -1
|
|
|
) -> None:
|
|
|
self.lr_scheduler = lr_scheduler
|
|
|
self.warmup_end_values = [pg["lr"] for pg in lr_scheduler.optimizer.param_groups]
|
|
|
self.last_step = last_step
|
|
|
self.warmup_duration = warmup_duration
|
|
|
self.step()
|
|
|
|
|
|
def state_dict(self):
|
|
|
return {k: v for k, v in self.__dict__.items() if k != "lr_scheduler"}
|
|
|
|
|
|
def load_state_dict(self, state_dict):
|
|
|
self.__dict__.update(state_dict)
|
|
|
|
|
|
def get_warmup_factor(self, step, **kwargs):
|
|
|
raise NotImplementedError
|
|
|
|
|
|
def step(
|
|
|
self,
|
|
|
):
|
|
|
self.last_step += 1
|
|
|
if self.last_step >= self.warmup_duration:
|
|
|
return
|
|
|
factor = self.get_warmup_factor(self.last_step)
|
|
|
for i, pg in enumerate(self.lr_scheduler.optimizer.param_groups):
|
|
|
pg["lr"] = factor * self.warmup_end_values[i]
|
|
|
|
|
|
def finished(
|
|
|
self,
|
|
|
):
|
|
|
if self.last_step >= self.warmup_duration:
|
|
|
return True
|
|
|
return False
|
|
|
|
|
|
|
|
|
@register()
|
|
|
class LinearWarmup(Warmup):
|
|
|
def __init__(
|
|
|
self, lr_scheduler: LRScheduler, warmup_duration: int, last_step: int = -1
|
|
|
) -> None:
|
|
|
super().__init__(lr_scheduler, warmup_duration, last_step)
|
|
|
|
|
|
def get_warmup_factor(self, step):
|
|
|
return min(1.0, (step + 1) / self.warmup_duration)
|
|
|
|
