Class CosineWeightDecayScheduler

utils/general_utils.py:195–240

class CosineWeightDecayScheduler(LRScheduler):
    def __init__(
        self,
        optimizer,
        max_iters: int,
        initial_wd: float = 0.05,
        final_wd: float = 0.20,
        last_epoch: int = -1,
    ):
        """
        对 weight_decay 进行余弦“增加”调度，从 initial_wd -> final_wd。
        
        参数:
        - optimizer: 任何带有 'weight_decay' param_group 的 optimizer
        - max_iters: 总的调度步数
        - initial_wd: 第 0 步时的 weight_decay
        - final_wd: 第 max_iters 步时的 weight_decay
        - last_epoch: 如需从中途恢复训练，传入上次迭代 idx
        """
        self.max_iters = max_iters
        self.initial_wd = initial_wd
        self.final_wd = final_wd
        super().__init__(optimizer, last_epoch)

    def get_lr(self):
        step = self._step_count
        # 限制在 [0, max_iters]
        if step <= 0:
            factor = 0.0
        elif step >= self.max_iters:
            factor = 1.0
        else:
            # factor 从 0 -> 1，按 1 - cos(pi * t / T) / 2
            factor = 0.5 * (1 - math.cos(math.pi * step / self.max_iters))

        wd = self.initial_wd + factor * (self.final_wd - self.initial_wd)
        return [wd for _ in self.optimizer.param_groups]

    def step(self, epoch=None):
        # 先让父类更新 last_epoch
        self._step_count += 1
        # super().step(epoch)
        # 再把新的 weight_decay 写到 optimizer
        new_wd = self.get_lr()
        for group, wd in zip(self.optimizer.param_groups, new_wd):
            group['weight_decay'] = wd