写一个例子，以防后面每次用到都忘记
class LR_Scheduler(object):
def __init__(self, optimizer, warmup_epochs, warmup_lr, num_epochs, base_lr, final_lr, iter_per_epoch,
constant_predictor_lr=False):
self.base_lr = base_lr
self.constant_predictor_lr = constant_predictor_lr
warmup_iter = iter_per_epoch * warmup_epochs
warmup_lr_schedule = np.linspace(warmup_lr, base_lr, warmup_iter)
decay_iter = iter_per_epoch * (num_epochs - warmup_epochs)
cosine_lr_schedule = final_lr + 0.5 * (base_lr - final_lr) * (
1 + np.cos(np.pi * np.arange(decay_iter) / decay_iter))
self.lr_schedule = np.concatenate((warmup_lr_schedule, cosine_lr_schedule))
self.optimizer = optimizer
self.iter = 0
self.current_lr = 0
def step(self):
for param_group in self.optimizer.param_groups:
if self.constant_predictor_lr and param_group['name'] == 'predictor':
param_group['lr'] = self.base_lr
else:
lr = param_group['lr'] = self.lr_schedule[self.iter]
self.iter += 1
self.current_lr = lr
return lr
def get_lr(self):
return self.current_lr
scheduler = LR_Scheduler(optimizer, warmup_epochs=300, warmup_lr=2e-2, num_epochs=2000, base_lr=3e-2, final_lr=2e-3, iter_per_epoch=1)
lr = []
for i in range(2000):
lr.append(scheduler.step())
plt.plot(lr)
plt.show()
