class CosineScheduler: def __init__(self, max_update, base_lr=0.01, final_lr=0, warmup_steps=0, warmup_begin_lr=0): self.base_lr_orig = base_lr self.max_update = max_update self.final_lr = final_lr self.warmup_steps = warmup_steps self.warmup_begin_lr = warmup_begin_lr self.max_steps = self.max_update - self.warmup_steps def get_warmup_lr(self, epoch): increase = (self.base_lr_orig - self.warmup_begin_lr) \ * float(epoch) / float(self.warmup_steps) return self.warmup_begin_lr + increase def __call__(self, epoch): if epoch < self.warmup_steps: return self.get_warmup_lr(epoch) if epoch <= self.max_update: self.base_lr = self.final_lr + ( self.base_lr_orig - self.final_lr) * (1 + math.cos( math.pi * (epoch - self.warmup_steps) / self.max_steps)) / 2 return self.base_lr scheduler = CosineScheduler(max_update=20, base_lr=0.3, final_lr=0.01) d2l.plot(torch.arange(num_epochs), [scheduler(t) for t in range(num_epochs)])

这段代码实现了一个余弦学习率调度程序，可以在训练神经网络时调整学习率。它包括一个 CosineScheduler 类和一个调用方法。在调用方法中，根据给定的 epoch 值，如果 epoch 值小于预热步数 warmup_steps，则返回预热学习率，否则返回余弦学习率。余弦学习率的计算是通过余弦函数实现的，其中初始学习率 base_lr 逐渐减小到最终学习率 final_lr，然后再逐渐增加回去。该程序还可以设置预热开始学习率 warmup_begin_lr 和预热步数 warmup_steps，并且可以根据最大迭代次数 max_update 计算出最大步数 max_steps。最后，该程序可以通过调用 plot 方法来可视化学习率的变化情况。