fixed typos

scheduler.py

# Copyright 2023 Ehsan Ahmadi
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
# associated documentation files (the “Software”), to deal in the Software without restriction, 
# including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all copies or substantial
# portions of the Software.

# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

# https://opensource.org/license/mit/

import abc
import numpy as np
from matplotlib import pyplot as plt


class Scheduler(abc.ABC):
    def __init__(self, num_epochs, steps_per_epoch, start_rate, **kwargs):
        self.num_epochs = num_epochs
        self.steps_per_epoch = steps_per_epoch
        self.start_rate = start_rate
    
    @abc.abstractmethod
    def get_rate(self, epoch: int, step: int) -> float:
        return 1.0
    
    def plot(self, path=None, title=None, ylabel=None, fig=None, ax=None, type_="png", **kwargs):
        if fig is None:
            fig = plt.figure()
            ax = fig.add_subplot(111)
        elif ax is None:
            ax = fig.axes[0]

        rates = np.zeros(self.num_epochs * self.steps_per_epoch)
        for epoch in range(self.num_epochs):
            for step in range(self.steps_per_epoch):
                rates[epoch*self.steps_per_epoch + step] = \
                    self.get_rate(epoch, step)
            
        ax.plot(rates, **kwargs)
        
        if title is None:
            title = self.__class__.__name__
        ax.set_title(title)
        ax.set_xlabel("Steps")
        if ylabel is not None:
            ax.set_ylabel(ylabel)
        if path is not None:
            fig.savefig(f"{path}/{title}.{type_.lower()}", dpi=600)
        return fig, ax


class PiecewiseLinearScheduler(Scheduler):
    def __init__(self, num_epochs, steps_per_epoch, rates, boundaries):
        super().__init__(num_epochs, steps_per_epoch, rates[0])
        self.rates = rates
        self.boundaries = boundaries
        assert len(rates) == len(boundaries) + 1, "rates and boundaries must have compatible lengths"
        self.validate_boundaries()
    
    def get_rate(self, epoch: int, step: int):
        epoch_float = epoch + step / self.steps_per_epoch
        for i, boundary in enumerate(self.boundaries):
            if epoch_float < boundary:
                return self.rates[i]
        return self.rates[-1]
    
    def validate_boundaries(self, ascending=False, descending=False):
        for i in range(len(self.boundaries)-1):
            assert self.boundaries[i] < self.boundaries[i+1], "boundaries must be increasing"

class CosineDecayScheduler(Scheduler):
    def __init__(self, num_epochs, steps_per_epoch, start_rate, end_rate, warmup_steps=0):
        '''If warmup_steps > 0, the learning rate will linearly increase from zero to start_rate
         then it will decrease according to the cosine decay schedule.'''
        super().__init__(num_epochs, steps_per_epoch, start_rate)
        self.end_rate = end_rate
        self.warmup_steps = warmup_steps
    
    def get_rate(self, epoch: int, step: int):
        epoch_float = epoch + step / self.steps_per_epoch
        warmup_epoch = self.warmup_steps / self.steps_per_epoch
        if epoch_float < warmup_epoch:
            return self.start_rate * epoch_float / warmup_epoch
        else:
            return self.end_rate + 0.5 * (self.start_rate - self.end_rate) * \
                (1 + np.cos(np.pi * (epoch_float - warmup_epoch) / (self.num_epochs - warmup_epoch))) 
        
class LinearDecayScheduler(Scheduler):
    def __init__(self, num_epochs, steps_per_epoch, start_rate, end_rate, warmup_steps=0):
        '''If warmup_steps > 0, the learning rate will linearly increase from zero to start_rate
         then it will decrease according to the linear decay schedule.'''
        super().__init__(num_epochs, steps_per_epoch, start_rate)
        self.end_rate = end_rate
        self.warmup_steps = warmup_steps
    
    def get_rate(self, epoch: int, step: int):
        epoch_float = epoch + step / self.steps_per_epoch
        warmup_epoch = self.warmup_steps / self.steps_per_epoch
        if epoch_float < warmup_epoch:
            return self.start_rate * epoch_float / warmup_epoch
        else:
            return self.start_rate - (self.start_rate - self.end_rate) * \
                (epoch_float - warmup_epoch) / (self.num_epochs - warmup_epoch)

class ExponentialDecayScheduler(Scheduler):
    def __init__(self, num_epochs, steps_per_epoch, start_rate=0.0, end_rate=1.0, warmup_steps=0):
        super().__init__(num_epochs, steps_per_epoch, start_rate)
        self.end_rate = end_rate
        self.warmup_steps = warmup_steps
    
    def get_rate(self, epoch: int, step: int):
        epoch_float = epoch + step / self.steps_per_epoch
        warmup_epoch = self.warmup_steps / self.steps_per_epoch
        if epoch_float < warmup_epoch:
            return self.start_rate * epoch_float / warmup_epoch
        else:
            return self.start_rate * np.exp(np.log(self.end_rate / self.start_rate) * \
                (epoch_float - warmup_epoch) / (self.num_epochs - warmup_epoch))
    
    
class LearningRateRangeFinderScheduler(Scheduler):
    def __init__(self, num_epochs, steps_per_epoch, start_rate=1e-5, epochs_per_cycle=1, gamma=10):
        ''' The learning rate will increase from start_rate to end_rate according to the exponential increase schedule.'''
        super().__init__(num_epochs, steps_per_epoch, start_rate)
        self.losses = np.zeros(self.num_epochs * self.steps_per_epoch)
        self.epochs_per_cycle = epochs_per_cycle
        self.gamma = gamma


    def get_rate(self, epoch: int, step: int):
        epoch_float = epoch + step / self.steps_per_epoch
        return self.start_rate * self.gamma ** (epoch_float / self.epochs_per_cycle) 

    def set_loss(self, epoch: int, step: int, loss: float):
        self.losses[epoch*self.steps_per_epoch + step] = loss
    
    def plot_loss(self, path=None, title=None, ylabel=None, fig=None, ax=None, type_="png", **kwargs):
        # ax1: learning rate vs. epoch
        # ax2: loss vs. learning rate
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(8, 4)) # type: ignore
        self.plot(fig=fig, ax=ax1, **kwargs)
        rates = np.zeros(self.num_epochs * self.steps_per_epoch)
        for epoch in range(self.num_epochs):
            for step in range(self.steps_per_epoch):
                rates[epoch*self.steps_per_epoch + step] = \
                    self.get_rate(epoch, step)
        ax2.semilogx(rates, self.losses)
        ax2.set_xlabel("rate")
        ax1.set_yscale('log')

        # finds the maximum loss and the corresponding learning rate
        max_loss_idx = np.argmax(self.losses)
        ax2.axvline(x=rates[max_loss_idx], color='r', linestyle='--')
        max_loss_learning_rate = rates[max_loss_idx]
        ax2.text(max_loss_learning_rate, np.min(self.losses),
                 f"lr_max_loss: {max_loss_learning_rate:.2e}",
                 horizontalalignment='right', verticalalignment='bottom', \
                 fontsize=8, rotation=90)
        
        # finds the steepest (positive slope) point of the smooted 
        kernel_radius = 50
        smoothed_losses = np.convolve(self.losses, np.ones(kernel_radius*2)/kernel_radius, mode='valid')
        ax2.semilogx(rates[kernel_radius:-kernel_radius+1], smoothed_losses, linestyle='-.')
        idx_steepest = np.argmax(smoothed_losses[1:] - smoothed_losses[:-1])
        suggested_lr = rates[idx_steepest + kernel_radius]
        ax2.axvline(x=suggested_lr, color='g', linestyle='--')
        ax2.text(suggested_lr, np.min(smoothed_losses),
                f"lr_steep_loss: {suggested_lr:.2e}",
                 horizontalalignment='right', verticalalignment='bottom', \
                 fontsize=8, rotation=90)
    
        if path is not None:
            fig.savefig(f"{path}/{title}.{type_.lower()}", dpi=600)
        return fig, ax

if __name__ == "__main__":
    sch = LinearDecayScheduler(100, 100, 1.0, 0.1, warmup_steps=1000)
    sch.plot(path=".", ylabel="rate")

    sch = PiecewiseLinearScheduler(100, 100, [1.0, 0.7, 0.1], [30, 60])
    sch.plot(path=".", title="piecewise_linear_scheduler", ylabel="rate")

    sch = CosineDecayScheduler(100, 100, 1.0, 0.1, warmup_steps=1000)
    sch.plot(path=".", title="cosine_decay_scheduler", ylabel="rate")

    sch = ExponentialDecayScheduler(100, 100, 1.0, 0.1, warmup_steps=1000)
    sch.plot(path=".", title="exponential_decay_scheduler", ylabel="rate")

    sch = LearningRateRangeFinderScheduler(100, 100, epochs_per_cycle=20)
    for epoch in range(100):
        for step in range(100):
            sch.set_loss(epoch, step, np.sin((epoch + step / 100) * np.pi / 100))
    sch.plot_loss(path=".", title="learning_rate_range_finder_scheduler", ylabel="loss")