Class AdamW

imperative/python/megengine/optimizer/adamw.py:11–127

r"""Implements the AdamW algorithm proposed in `"Decoupled Weight Decay Regularization" `_. This optimizer combines the Adam optimizer with weight decay regularization to prevent overfitting. Args: params (Union[Iterable[Parameter], dict]): Iterabl

class AdamW(Optimizer):
    r"""Implements the AdamW algorithm proposed in `"Decoupled Weight Decay Regularization" <https://arxiv.org/abs/1711.05101>`_.

    This optimizer combines the Adam optimizer with weight decay regularization to prevent overfitting.

    Args:
        params (Union[Iterable[Parameter], dict]): Iterable of parameters to optimize or dicts defining
            parameter groups.
        lr (float): Learning rate.
        betas (Tuple[float, float], optional): Coefficients used for computing running averages of gradient
            and its square. Default: (0.9, 0.999).
        eps (float, optional): Term added to the denominator to improve numerical stability. Default: 1e-8.
        weight_decay (float, optional): Weight decay (L2 penalty). Default: 1e-2.
    
    Returns:
        An instance of the AdamW optimizer.
    """

    def __init__(
        self,
        params: Union[Iterable[Parameter], dict],
        lr: float,
        betas: Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 1e-2,
    ):
        if lr < 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if weight_decay < 0.0:
            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
        if not 0.0 <= betas[0] < 1.0:
            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
        if not 0.0 <= betas[1] < 1.0:
            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))

        defaults = dict(lr=lr, weight_decay=weight_decay, betas=betas, eps=eps)
        super().__init__(params, defaults)
        self._disable_type_convert = True

    def _create_state(self, param_group):
        for param in param_group["params"]:
            self._add_state(param, "exp_avg")
            self._add_state(param, "exp_avg_sq")
            self._add_state(param, "step", initializer=0.0)

    def _updates(self, param_group):
        lr = param_group["lr"]
        weight_decay = param_group["weight_decay"]
        eps = param_group["eps"]
        beta0, beta1 = param_group["betas"]

        def make_scalar(val):
            return tensor(val, dtype="float32")

        # since `conver_inputs` is disabled for param updates,
        # scalar should be explicitly tansforred to tensor

        _lr, _neg_lr = map(make_scalar, (lr, -lr))