Class SGD

python/tvm/relax/training/optimizer.py:244–353

Implements stochastic gradient descent. The returned function of `get_function()` is equivalent to the following numpy code: .. code-block:: python def SGD(param_tuple, grad_tuple, state_tuple): num_steps = state_tuple[0] param_tuple_new, state_tuple_new =

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

    The returned function of `get_function()` is equivalent to the following numpy code:

    .. code-block:: python

        def SGD(param_tuple, grad_tuple, state_tuple):
            num_steps = state_tuple[0]
            param_tuple_new, state_tuple_new = [], []
            state_tuple_new.append(num_steps + 1)
            for i in range(len(param_tuple)):
                param = param_tuple[i]
                grad = grad_tuple[i]
                param_tuple_new.append(param - lr * (grad + weight_decay * param))
            return param_tuple_new, state_tuple_new

    Parameters
    ----------
    lr : float
        learning rate

    weight_decay : float
        weight decay (L2 penalty) (default: 0)
    """

    def __init__(self, lr: float, weight_decay: float = 0) -> None:
        super().__init__("SGD")
        self.lr = float(lr)
        self.weight_decay = float(weight_decay)

    def init(self, params: Var | list[Var]) -> "SGD":
        """Set the parameters, determine the dtype, and construct the initial state for the
        optimizer.

        The state of SGD is `(num_steps,)`.

        Parameters
        ----------
        params : Union[Var, List[Var]]
            The parameter or the list of parameters to optimize.

            Parameters should all be Vars of floating point Tensors, including float32, float64,
            float16, etc. Currently, all parameters should have the same dtype, and that dtype
            will be used as the dtype of the optimizer states.

        Returns
        -------
        self : SGD
            The SGD optimizer itself.
        """
        if not isinstance(params, list):
            params = [params]
        self._set_params_and_dtype(params)
        self.state = (
            # num_steps = 0
            tvm.runtime.tensor(np.zeros((), "int64")),
        )