solver.py

import logging

from torch.optim import SGD, Adam
from torch.optim.lr_scheduler import LambdaLR, StepLR


class LambdaStepLR(LambdaLR):
    def __init__(self, optimizer, lr_lambda, last_step=-1):
        super(LambdaStepLR, self).__init__(optimizer, lr_lambda, last_step)

    @property
    def last_step(self):
        """Use last_epoch for the step counter"""
        return self.last_epoch

    @last_step.setter
    def last_step(self, v):
        self.last_epoch = v


class PolyLR(LambdaStepLR):
    """DeepLab learning rate policy"""

    def __init__(self, optimizer, max_iter, power=0.9, last_step=-1):
        super(PolyLR, self).__init__(optimizer, lambda s: (1 - s / (max_iter + 1)) ** power, last_step)


class SquaredLR(LambdaStepLR):
    """Used for SGD Lars"""

    def __init__(self, optimizer, max_iter, last_step=-1):
        super(SquaredLR, self).__init__(optimizer, lambda s: (1 - s / (max_iter + 1)) ** 2, last_step)


class ExpLR(LambdaStepLR):
    def __init__(self, optimizer, step_size, gamma=0.9, last_step=-1):
        # (0.9 ** 21.854) = 0.1, (0.95 ** 44.8906) = 0.1
        # To get 0.1 every N using gamma 0.9, N * log(0.9)/log(0.1) = 0.04575749 N
        # To get 0.1 every N using gamma g, g ** N = 0.1 -> N * log(g) = log(0.1) -> g = np.exp(log(0.1) / N)
        super(ExpLR, self).__init__(optimizer, lambda s: gamma ** (s / step_size), last_step)


def initialize_optimizer(params, config):
    assert config.optimizer in ["SGD", "Adagrad", "Adam", "RMSProp", "Rprop", "SGDLars"]

    if config.optimizer == "SGD":
        return SGD(
            params,
            lr=config.lr,
            momentum=config.sgd_momentum,
            dampening=config.sgd_dampening,
            weight_decay=config.weight_decay,
        )
    elif config.optimizer == "Adam":
        return Adam(
            params, lr=config.lr, betas=(config.adam_beta1, config.adam_beta2), weight_decay=config.weight_decay
        )
    else:
        logging.error("Optimizer type not supported")
        raise ValueError("Optimizer type not supported")


def initialize_scheduler(optimizer, config, last_step=-1):
    if config.scheduler == "StepLR":
        return StepLR(optimizer, step_size=config.step_size, gamma=config.step_gamma, last_epoch=last_step)
    elif config.scheduler == "PolyLR":
        return PolyLR(optimizer, max_iter=config.max_iter, power=config.poly_power, last_step=last_step)
    elif config.scheduler == "SquaredLR":
        return SquaredLR(optimizer, max_iter=config.max_iter, last_step=last_step)
    elif config.scheduler == "ExpLR":
        return ExpLR(optimizer, step_size=config.exp_step_size, gamma=config.exp_gamma, last_step=last_step)
    else:
        logging.error("Scheduler not supported")