From 6073f8cbd0eb924b4b265483af63f3ed6c03d10d Mon Sep 17 00:00:00 2001 From: Less Wright Date: Sun, 16 May 2021 13:07:49 -0700 Subject: [PATCH] fix off by 1 iteration on linear warmup, adds show_schedule() to show lr quickly after training --- ranger21/ranger21.py | 69 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 53 insertions(+), 16 deletions(-) diff --git a/ranger21/ranger21.py b/ranger21/ranger21.py index 1a529b6..48fb08b 100644 --- a/ranger21/ranger21.py +++ b/ranger21/ranger21.py @@ -39,6 +39,9 @@ import torch.nn.functional as F import math import collections +# this is to support showing the lr curves after a training run. +import matplotlib.pyplot as plt + import copy from torch import linalg as LA @@ -196,7 +199,7 @@ class Ranger21(TO.Optimizer): ) # default untuned linear warmup beta_pct = beta_warmup_iters / self.total_iterations - #print(f"beta_warmup_pct = {beta_pct}") + # print(f"beta_warmup_pct = {beta_pct}") # this can be unreasonable for short runs...so let's compare vs warmup pct % of total epochs if beta_pct > 0.45: @@ -222,6 +225,7 @@ class Ranger21(TO.Optimizer): self.total_iterations - self.start_warm_down ) self.warmdown_displayed = False # print when warmdown begins... + self.warmup_curr_pct = 0.01 # used to verify warmup reaches full set point. """ print(f"debug warmdown:\n") @@ -270,13 +274,30 @@ class Ranger21(TO.Optimizer): super().__setstate__(state) # show settings at init or if called + def show_schedule(self): + if not self.tracking_lr: + print( + "No data from training yet. Please train and then use this to show the lr curves" + ) + return + x = self.tracking_lr + plt.plot(x) + maxlr = max(x) + minlr = min(x) + startlr = x[0] + plt.title( + f"Ranger21 learning rate schedule\nStart={startlr}\nMax ={maxlr}\n,Min={minlr}\n" + ) + plt.show() def show_settings(self): print(f"Ranger21 optimizer ready with following settings:\n") print(f"Core optimizer = {self.core_engine}") print(f"Learning rate of {self.starting_lr}\n") - print(f"Important - num_epochs of training = ** {self.num_epochs} epochs **\nplease confirm this is correct or warmup and warmdown will be off\n") + print( + f"Important - num_epochs of training = ** {self.num_epochs} epochs **\nplease confirm this is correct or warmup and warmdown will be off\n" + ) if self.use_adabelief: print(f"using AdaBelief for variance computation") @@ -392,43 +413,59 @@ class Ranger21(TO.Optimizer): if style is None: return lr - if step >= warmup: + if step > warmup: if not self.warmup_complete: + + if not self.warmup_curr_pct == 1.0: + print( + f"Error - lr did not achieve full set point from warmup, currently {self.warmup_curr_pct}" + ) + self.warmup_complete = True + print(f"\n** Ranger21 update = Warmup complete - lr set to {lr}\n") return lr if style == "linear": - new_lr = lr * min(1.0, (step / warmup)) + self.warmup_curr_pct = min(1.0, (step / warmup)) + new_lr = lr * self.warmup_curr_pct self.current_lr = new_lr return new_lr - elif style == "exponential": - return lr * (1.0 - math.exp(-step / warmup)) + # elif style == "exponential": + # return lr * (1.0 - math.exp(-step / warmup)) + else: raise ValueError(f"warmup type {style} not implemented.") - def get_warm_down(self, lr, iteration): """ linear style warmdown """ if iteration < self.start_warm_down: return lr - if iteration > self.start_warm_down-1: + if iteration > self.start_warm_down - 1: # print when starting if not self.warmdown_displayed: - print(f"\n--> Warmdown starting now....\n") + print( + f"\n** Ranger21 update: Warmdown starting now. Current iteration = {iteration}....\n" + ) self.warmdown_displayed = True - warmdown_iteration = iteration - self.start_warm_down-1 # to force the first iteration to be 1 instead of 0 - if warmdown_iteration <=1: + warmdown_iteration = ( + iteration + 1 + ) - self.start_warm_down # to force the first iteration to be 1 instead of 0 + + if warmdown_iteration < 1: + print( + f" warning - iteration started at {iteration} and {self.start_warm_down} with value {warmdown_iteration}" + ) warmdown_iteration = 1 - print(f" warning - iteration started at {iteration} and {self.start_warm_down} with value {warmdown_iteration}") - warmdown_iteration = 1 - print(f"warmdown iteration = {warmdown_iteration}") + # print(f"warmdown iteration = {warmdown_iteration}") # linear start 3672 5650 total iterations 1972 iterations - warmdown_pct = (warmdown_iteration / self.warmdown_total_iterations) - if warmdown_pct >1.00: + warmdown_pct = warmdown_iteration / ( + self.warmdown_total_iterations + 1 + ) # +1 to offset that we have to include first as an iteration to support 1 index instead of 0 based. + if warmdown_pct > 1.00: print(f"error in warmdown pct calc. new pct = {warmdown_pct}") print(f"auto handled but please report issue") warmdown_pct = 1.00