pytroch实现warmup代码

class myWarmUpLR(_LRScheduler):
def init(self, optimizer, total_iters, last_epoch=-1):
self.total_iters = total_iters
super().init(optimizer, last_epoch)
def get_lr(self):
return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]
warm_epoch = 5
warmup = myWarmUpLR(optimizer, iter_per_epoch *warm_epoch)