MegEngine/python_module/megengine/optimizer/internal.py

# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
from typing import Union

import megengine._internal as mgb

from ..core.tensor import Tensor, tensor


def add_update_fastpath(
    dest: Tensor,
    delta: Tensor,
    *,
    alpha: Union[Tensor, float, int] = 1.0,
    beta: Union[Tensor, float, int] = 1.0,
    bias: Union[Tensor, float, int] = 0.0
):
    """a fast-path ONLY used to update parameters in optimzier, since it
    would bypass computing graph and launch dnn/add_update kernel directly,
    it is more efficient than functional/add_update.
    """

    if isinstance(beta, Tensor) or isinstance(alpha, Tensor):
        delta *= beta
        beta = 1.0
    if isinstance(alpha, Tensor):
        delta += (alpha - 1.0) * dest
        alpha = 1.0
    if isinstance(bias, Tensor):
        delta += bias
        bias = 0.0

    if not isinstance(delta, Tensor):
        delta = tensor(delta, device=dest.device, dtype=dest.dtype)

    def get_v(x):
        if x._Tensor__val is None:
            assert isinstance(x._Tensor__sym, mgb.SymbolVar)
            return x._Tensor__sym.eager_val
        else:
            assert isinstance(x._Tensor__val, mgb.SharedND)
            return x._Tensor__val

    mgb.mgb._add_update_fastpath(get_v(dest), get_v(delta), alpha, beta, bias)
    return dest