Method _updates

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

(self, param_group)

            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))
        _weight_decay = make_scalar(weight_decay)
        _eps = make_scalar(eps)
        _beta0, _beta1 = map(make_scalar, (beta0, beta1))

        c1, c05 = map(make_scalar, (1.0, 0.5))

        inplace_mode = int(os.getenv("MEGENGINE_INPLACE_UPDATE", "0"))
        if inplace_mode:
            # reduce device sync
            c1_sub_beta0, c1_sub_beta1 = map(make_scalar, (1 - beta0, 1 - beta1))

        for param in param_group["params"]:

            if param.grad is None:
                continue

            grad = param.grad

            states = self._state[param]

            step, exp_avg, exp_avg_sq = (
                states["step"],
                states["exp_avg"],
                states["exp_avg_sq"],
            )

            if inplace_mode:
                _inplace_add_(step, c1, alpha=c1, beta=c1)
                _inplace_add_(exp_avg, grad, alpha=_beta0, beta=c1_sub_beta0)
                _inplace_add_(
                    exp_avg_sq, grad * grad, alpha=_beta1, beta=c1_sub_beta1,
                )

                delta = (exp_avg / (c1 - _beta0 ** step)) / (
                    (exp_avg_sq / (c1 - _beta1 ** step)) ** c05 + _eps
                )
                if is_tracing() or weight_decay != 0.0:
                    delta += param * _weight_decay
                _inplace_add_(param, delta, alpha=c1, beta=_neg_lr)
                continue

            # step = step + c1
            step += c1

            # exp_avg = _beta0 * exp_avg + grad * (c1 - _beta0)