| |
| |
|
|
| import numpy as np |
|
|
|
|
| class IterExponential: |
| def __init__(self, total_iter_length, final_ratio, warmup_steps=0) -> None: |
| """ |
| Customized iteration-wise exponential scheduler. |
| Re-calculate for every step, to reduce error accumulation |
| |
| Args: |
| total_iter_length (int): Expected total iteration number |
| final_ratio (float): Expected LR ratio at n_iter = total_iter_length |
| """ |
| self.total_length = total_iter_length |
| self.effective_length = total_iter_length - warmup_steps |
| self.final_ratio = final_ratio |
| self.warmup_steps = warmup_steps |
|
|
| def __call__(self, n_iter) -> float: |
| if n_iter < self.warmup_steps: |
| alpha = 1.0 * n_iter / self.warmup_steps |
| elif n_iter >= self.total_length: |
| alpha = self.final_ratio |
| else: |
| actual_iter = n_iter - self.warmup_steps |
| alpha = np.exp( |
| actual_iter / self.effective_length * np.log(self.final_ratio) |
| ) |
| return alpha |
|
|
|
|
| if "__main__" == __name__: |
| lr_scheduler = IterExponential( |
| total_iter_length=50000, final_ratio=0.01, warmup_steps=200 |
| ) |
| lr_scheduler = IterExponential( |
| total_iter_length=50000, final_ratio=0.01, warmup_steps=0 |
| ) |
|
|
| x = np.arange(100000) |
| alphas = [lr_scheduler(i) for i in x] |
| import matplotlib.pyplot as plt |
|
|
| plt.plot(alphas) |
| plt.savefig("lr_scheduler.png") |
|
|
