Class SGD

imperative/python/megengine/optimizer/sgd.py:13–114

r"""Implements stochastic gradient descent. This optimizer performs stochastic gradient descent with optional momentum and weight decay. Nesterov momentum is based on the formula from `"On the importance of initialization and momentum in deep learning" <http://www.cs.toronto.edu/%7Ehin

class SGD(Optimizer):
    r"""Implements stochastic gradient descent.

    This optimizer performs stochastic gradient descent with optional momentum and weight decay.

    Nesterov momentum is based on the formula from
    `"On the importance of initialization and momentum in deep learning" <http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf>`_.

    Args:
        params (Union[Iterable[Parameter], dict]): Iterable of parameters to optimize or dicts defining
            parameter groups.
        lr (float): Learning rate.
        momentum (float): Momentum factor. Default: 0.0.
        nesterov (bool): Enables Nesterov momentum. Default: False.
        weight_decay (float): Weight decay (L2 penalty). Default: 0.0.
    
    Returns:
        An instance of the SGD optimizer.

    Note:
        This optimizer does not guarantee that the interval does not include the stop value in cases
        where the step is not an integer and floating-point rounding errors affect the length of the output tensor.
    """

    def __init__(
        self,
        params: Union[Iterable[Parameter], dict],
        lr: float,
        momentum: float = 0.0,
        nesterov: bool = False,
        weight_decay: float = 0.0,
    ):
        assert lr >= 0.0, "Invalid learning rate: {}".format(lr)
        assert momentum >= 0.0, "Invalid momentum value: {}".format(momentum)
        assert weight_decay >= 0.0, "Invalid weight_decay value: {}".format(
            weight_decay
        )
        assert not nesterov or momentum > 0.0, "Nesterov momentum requires a momentum"

        defaults = dict(lr=lr, momentum=momentum, weight_decay=weight_decay)
        super().__init__(params, defaults)
        self.nesterov = nesterov
        self._disable_type_convert = True

    def _create_state(self, param_group):
        if param_group["momentum"] != 0.0:
            for param in param_group["params"]:
                self._add_state(param, "momentum_buffer")

    def _updates(self, param_group):
        lr = param_group["lr"]
        weight_decay = param_group["weight_decay"]
        momentum = param_group["momentum"]

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

        _lr = tensor(lr, dtype="float32")