mindspore 2.2.14__cp37-cp37m-manylinux1_x86_64.whl → 2.3.0rc2__cp37-cp37m-manylinux1_x86_64.whl

mindspore/dataset/vision/utils.py

@@ -127,7 +127,10 @@ class AutoAugmentPolicy(str, Enum):
                     AutoAugmentPolicy.CIFAR10: cde.AutoAugmentPolicy.DE_AUTO_AUGMENT_POLICY_CIFAR10,
                     AutoAugmentPolicy.SVHN: cde.AutoAugmentPolicy.DE_AUTO_AUGMENT_POLICY_SVHN}
 
-        return c_values.get(policy)
+        value = c_values.get(policy)
+        if value is None:
+            raise RuntimeError("Unsupported AutoAugmentPolicy, only support IMAGENET, CIFAR10, and SVHN.")
+        return value
 
 
 class Border(str, Enum):
@@ -160,7 +163,11 @@ class Border(str, Enum):
                          Border.EDGE: 'edge',
                          Border.REFLECT: 'reflect',
                          Border.SYMMETRIC: 'symmetric'}
-        return python_values.get(border_type)
+
+        value = python_values.get(border_type)
+        if value is None:
+            raise RuntimeError("Unsupported Border type, only support CONSTANT, EDGE, REFLECT and SYMMETRIC.")
+        return value
 
     @staticmethod
     def to_c_type(border_type):
@@ -172,7 +179,10 @@ class Border(str, Enum):
                     Border.REFLECT: cde.BorderType.DE_BORDER_REFLECT,
                     Border.SYMMETRIC: cde.BorderType.DE_BORDER_SYMMETRIC}
 
-        return c_values.get(border_type)
+        value = c_values.get(border_type)
+        if value is None:
+            raise RuntimeError("Unsupported Border type, only support CONSTANT, EDGE, REFLECT and SYMMETRIC.")
+        return value
 
 
 class ConvertMode(IntEnum):
@@ -250,7 +260,11 @@ class ConvertMode(IntEnum):
                     ConvertMode.COLOR_RGBA2GRAY: cde.ConvertMode.DE_COLOR_RGBA2GRAY,
                     }
 
-        return c_values.get(mode)
+        mode = c_values.get(mode)
+        if mode is None:
+            raise RuntimeError("Unsupported ConvertMode, see https://www.mindspore.cn/docs/zh-CN/master/api_python/"
+                               "dataset_vision/mindspore.dataset.vision.ConvertColor.html for more details.")
+        return mode
 
 
 class ImageBatchFormat(IntEnum):
@@ -273,7 +287,10 @@ class ImageBatchFormat(IntEnum):
         c_values = {ImageBatchFormat.NHWC: cde.ImageBatchFormat.DE_IMAGE_BATCH_FORMAT_NHWC,
                     ImageBatchFormat.NCHW: cde.ImageBatchFormat.DE_IMAGE_BATCH_FORMAT_NCHW}
 
-        return c_values.get(image_batch_format)
+        value = c_values.get(image_batch_format)
+        if value is None:
+            raise RuntimeError("Unsupported ImageBatchFormat, only support NHWC and NCHW.")
+        return value
 
 
 class ImageReadMode(IntEnum):
@@ -298,7 +315,11 @@ class ImageReadMode(IntEnum):
         c_values = {ImageReadMode.UNCHANGED: cde.ImageReadMode.DE_IMAGE_READ_MODE_UNCHANGED,
                     ImageReadMode.GRAYSCALE: cde.ImageReadMode.DE_IMAGE_READ_MODE_GRAYSCALE,
                     ImageReadMode.COLOR: cde.ImageReadMode.DE_IMAGE_READ_MODE_COLOR}
-        return c_values.get(image_read_mode)
+
+        value = c_values.get(image_read_mode)
+        if value is None:
+            raise RuntimeError("Unsupported ImageReadMode, only support UNCHANGED, GRAYSCALE and COLOR.")
+        return value
 
 
 class Inter(IntEnum):
@@ -334,7 +355,10 @@ class Inter(IntEnum):
                          Inter.LINEAR: LINEAR,
                          Inter.CUBIC: CUBIC}
 
-        return python_values.get(inter_type)
+        value = python_values.get(inter_type)
+        if value is None:
+            raise RuntimeError("Unsupported interpolation, only support NEAREST, ANTIALIAS, LINEAR and CUBIC.")
+        return value
 
     @staticmethod
     def to_c_type(inter_type):
@@ -347,7 +371,11 @@ class Inter(IntEnum):
                     Inter.AREA: cde.InterpolationMode.DE_INTER_AREA,
                     Inter.PILCUBIC: cde.InterpolationMode.DE_INTER_PILCUBIC}
 
-        return c_values.get(inter_type)
+        value = c_values.get(inter_type)
+        if value is None:
+            raise RuntimeError("Unsupported interpolation, only support NEAREST, LINEAR, CUBIC, AREA and PILCUBIC.")
+
+        return value
 
 
 class SliceMode(IntEnum):
@@ -370,7 +398,10 @@ class SliceMode(IntEnum):
         c_values = {SliceMode.PAD: cde.SliceMode.DE_SLICE_PAD,
                     SliceMode.DROP: cde.SliceMode.DE_SLICE_DROP}
 
-        return c_values.get(mode)
+        value = c_values.get(mode)
+        if value is None:
+            raise RuntimeError("Unsupported SliceMode, only support PAD and DROP.")
+        return value
 
 
 def encode_jpeg(image, quality=75):

mindspore/dataset/vision/validators.py

@@ -36,7 +36,7 @@ def check_affine(method):
         [degrees, translate, scale, shear, resample, fill_value], _ = parse_user_args(method, *args, **kwargs)
 
         type_check(degrees, (int, float), "degrees")
-        check_degrees(degrees)
+        check_value(degrees, [-180, 180], "degrees")
 
         type_check(translate, (list, tuple), "translate")
         if len(translate) != 2:
@@ -54,8 +54,13 @@ def check_affine(method):
                 raise TypeError("The length of shear should be 2.")
             for i, _ in enumerate(shear):
                 type_check(shear[i], (int, float), "shear[{}]".format(i))
+                check_value(shear[i], [-180, 180], "shear[{}]".format(i))
+        else:
+            check_value(shear, [-180, 180], "shear")
 
         type_check(resample, (Inter,), "resample")
+        if resample in [Inter.ANTIALIAS, Inter.PILCUBIC]:
+            raise RuntimeError("Unsupported interpolation, only NEAREST, LINEAR, CUBIC and AREA.")
 
         check_fill_value(fill_value)
 
@@ -182,6 +187,7 @@ def check_padding(padding):
     """Parsing the padding arguments and check if it is legal."""
     type_check(padding, (tuple, list, numbers.Number), "padding")
     if isinstance(padding, numbers.Number):
+        type_check(padding, (int,), "padding")
         check_value(padding, (0, INT32_MAX), "padding")
     if isinstance(padding, (tuple, list)):
         if len(padding) not in (2, 4):

mindspore/experimental/optim/__init__.py

@@ -15,7 +15,7 @@
 """
 Optimizer.
 
-Provide common optimizers for training, such as SGD, ADAM, Momentum.
+Provide common optimizers for training, such as SGD, Adam.
 The optimizer is used to calculate and update the gradients.
 """
 from __future__ import absolute_import
@@ -24,7 +24,17 @@ from mindspore.experimental.optim.optimizer import Optimizer
 from mindspore.experimental.optim.adamw import AdamW
 from mindspore.experimental.optim.sgd import SGD
 from mindspore.experimental.optim.adam import Adam
+from mindspore.experimental.optim.adadelta import Adadelta
+from mindspore.experimental.optim.adagrad import Adagrad
+from mindspore.experimental.optim.asgd import ASGD
+from mindspore.experimental.optim.adamax import Adamax
+from mindspore.experimental.optim.rmsprop import RMSprop
+from mindspore.experimental.optim.rprop import Rprop
+from mindspore.experimental.optim.radam import RAdam
+from mindspore.experimental.optim.nadam import NAdam
+
 from mindspore.experimental.optim import lr_scheduler
 
 
-__all__ = ['Optimizer', 'AdamW', 'SGD', 'Adam']
+__all__ = ['Optimizer', 'AdamW', 'SGD', 'Adam', 'Adadelta', 'Adagrad',
+           'ASGD', 'Adamax', 'RMSprop', 'Rprop', 'RAdam', 'NAdam']

mindspore/experimental/optim/adadelta.py

@@ -0,0 +1,161 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""adadelta"""
+from __future__ import absolute_import
+
+from mindspore.ops import functional as F, composite as C, operations as P
+import mindspore.common.dtype as mstype
+from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
+from mindspore import _checkparam as validator
+from mindspore import jit
+
+_adadelta_opt = C.MultitypeFuncGraph("adadelta_opt")
+
+
+@_adadelta_opt.register("Function", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
+def _tensor_run_opt(opt, rho, epsilon, learning_rate, weight, accum, accum_update, gradient):
+    """Apply adadelta optimizer to the weight parameter."""
+    success = True
+    success = F.depend(success, opt(weight, accum, accum_update, learning_rate, rho, epsilon, gradient))
+    return success
+
+
+class Adadelta(Optimizer):
+    r"""
+    Implements Adadelta algorithm.
+
+    .. math::
+       \begin{aligned}
+            &\rule{150mm}{0.4pt}                                                                 \\
+            &\textbf{input}      : \gamma \text{ (lr)}, \: \theta_0 \text{ (params)},
+                \: f(\theta) \text{ (objective)}, \: \rho \text{ (decay)},
+                \: \lambda \text{ (weight decay)}                                                \\
+            &\textbf{initialize} :  v_0  \leftarrow 0 \: \text{ (square avg)},
+                \: u_0 \leftarrow 0 \: \text{ (accumulate variables)}                     \\[-1.ex]
+            &\rule{110mm}{0.4pt}                                                                 \\
+            &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
+            &\hspace{5mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})           \\
+            &\hspace{5mm}if \: \lambda \neq 0                                                    \\
+            &\hspace{10mm} g_t \leftarrow g_t + \lambda  \theta_{t-1}                            \\
+            &\hspace{5mm} v_t      \leftarrow v_{t-1} \rho + g^2_t (1 - \rho)                    \\
+            &\hspace{5mm}\Delta x_t    \leftarrow   \frac{\sqrt{u_{t-1} +
+                \epsilon }}{ \sqrt{v_t + \epsilon}  }g_t \hspace{21mm}                           \\
+            &\hspace{5mm} u_t  \leftarrow   u_{t-1}  \rho +
+                 \Delta x^2_t  (1 - \rho)                                                        \\
+            &\hspace{5mm}\theta_t      \leftarrow   \theta_{t-1} - \gamma  \Delta x_t            \\
+            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+            &\bf{return} \:  \theta_t                                                     \\[-1.ex]
+            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+       \end{aligned}
+
+    .. warning::
+        This is an experimental optimizer API that is subject to change.
+        This module must be used with lr scheduler module in `LRScheduler Class
+        <https://www.mindspore.cn/docs/en/master/api_python/mindspore.experimental.html#lrscheduler-class>`_ .
+
+    Args:
+        params (Union[list(Parameter), list(dict)]): list of parameters to optimize or dicts defining
+            parameter groups.
+        lr (Union[int, float, Tensor], optional): learning rate. Default: ``1.0``.
+        rho (float, optional): coefficient used for computing a running average
+            of squared gradients. :math:`\rho` in the formula above. Default: ``0.9``.
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. :math:`\epsilon` in the formula above. Default: ``1e-6``.
+        weight_decay (float, optional): weight decay (L2 penalty). Default: ``0.``.
+
+    Keyword Args:
+        maximize (bool, optional): maximize the params based on the objective, instead of minimizing.
+            Default: ``False``.
+
+    Inputs:
+        - **gradients** (tuple[Tensor]) - The gradients of `params`.
+
+    Raises:
+        ValueError: If the learning rate is not int, float or Tensor.
+        ValueError: If the learning rate is less than 0.
+        ValueError: If the `eps` is less than or equal to 0.0.
+        ValueError: If the `rho` is not in the range of [0, 1].
+        ValueError: If the `weight_decay` is less than 0.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import nn
+        >>> from mindspore.experimental import optim
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
+        >>> optimizer = optim.Adadelta(net.trainable_params(), lr=0.1)
+        >>> def forward_fn(data, label):
+        ...     logits = net(data)
+        ...     loss = loss_fn(logits, label)
+        ...     return loss, logits
+        >>> grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)
+        >>> def train_step(data, label):
+        ...     (loss, _), grads = grad_fn(data, label)
+        ...     optimizer(grads)
+        ...     return loss
+    """
+
+    def __init__(self, params, lr=1.0, rho=0.9, eps=1e-6, weight_decay=0.0, *, maximize=False):
+        check_not_less_than_without_equal(lr, "lr", self.cls_name)
+        check_not_less_than_without_equal(eps, "eps", self.cls_name)
+        check_not_less_than(weight_decay, "weight_decay", self.cls_name)
+        validator.check_float_range(rho, 0., 1., validator.INC_BOTH, "rho", self.cls_name)
+
+        defaults = dict(
+            lr=lr,
+            rho=rho,
+            eps=eps,
+            weight_decay=weight_decay,
+            maximize=maximize,
+        )
+        super(Adadelta, self).__init__(params, defaults)
+
+        self.accum = self.parameters.clone(prefix="accum", init=0)
+        self.accum_update = self.parameters.clone(prefix="accum_update", init=0)
+        self.opt = P.ApplyAdadelta()
+        self.op_cast = P.Cast()
+
+    @jit
+    def implementation(self, lr, rho, eps, maximize, weight_decay, start_id, end_id, gradients):
+        """Extract the common computing part for acceleration"""
+        params = self.parameters[start_id: end_id]
+        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = self._decay_weight(weight_decay, params, grads)
+        accum = self.accum[start_id: end_id]
+        accum_update = self.accum_update[start_id: end_id]
+        self.hyper_map(F.partial(_adadelta_opt, self.opt, rho, eps, lr),
+                       params, accum, accum_update, grads)
+        return True
+
+    def construct(self, gradients):
+        for group_id, group in enumerate(self.param_groups):
+            lr = self.lrs[group_id]
+            if isinstance(group.get("lr"), float):
+                lr = self.op_cast(group.get("lr"), mstype.float32)
+            maximize = group.get("maximize")
+            rho = group["rho"]
+            eps = group["eps"]
+
+            start_id = self.group_start_id[group_id]
+            end_id = self.group_start_id[group_id + 1]
+            weight_decay = group["weight_decay"]
+            self.implementation(lr, rho, eps, maximize, weight_decay, start_id, end_id, gradients)
+
+        return True

mindspore/experimental/optim/adagrad.py

@@ -0,0 +1,168 @@
+# Copyright 2023 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""adagrad"""
+from __future__ import absolute_import
+
+from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore.common import Tensor, Parameter
+import mindspore.common.dtype as mstype
+from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
+from mindspore import jit
+
+_adagrad_opt = C.MultitypeFuncGraph("adagrad_opt")
+
+
+@_adagrad_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor")
+def _tensor_run_opt(opt, learning_rate, weight, accum, gradient):
+    """Apply adagrad optimizer to the weight parameter."""
+    success = True
+    success = F.depend(success, opt(weight, accum, learning_rate, gradient))
+    return success
+
+
+class Adagrad(Optimizer):
+    r"""
+    Implements Adagrad algorithm.
+
+    .. math::
+       \begin{aligned}
+            &\rule{110mm}{0.4pt}                                                                 \\
+            &\textbf{input}      : \gamma \text{ (lr)}, \: \theta_0 \text{ (params)}, \: f(\theta)
+                \text{ (objective)}, \: \lambda \text{ (weight decay)},                          \\
+            &\hspace{12mm}    \tau \text{ (initial accumulator value)}, \: \eta\text{ (lr decay)}\\
+            &\textbf{initialize} :  state\_sum_0 \leftarrow 0                             \\[-1.ex]
+            &\rule{110mm}{0.4pt}                                                                 \\
+            &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
+            &\hspace{5mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})           \\
+            &\hspace{5mm} \tilde{\gamma}    \leftarrow \gamma / (1 +(t-1) \eta)                  \\
+            &\hspace{5mm} \textbf{if} \: \lambda \neq 0                                          \\
+            &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1}                             \\
+            &\hspace{5mm}state\_sum_t  \leftarrow  state\_sum_{t-1} + g^2_t                      \\
+            &\hspace{5mm}\theta_t \leftarrow
+                \theta_{t-1}- \tilde{\gamma} \frac{g_t}{\sqrt{state\_sum_t}+\epsilon}            \\
+            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+            &\bf{return} \:  \theta_t                                                     \\[-1.ex]
+            &\rule{110mm}{0.4pt}                                                          \\[-1.ex]
+       \end{aligned}
+
+    .. warning::
+        This is an experimental optimizer API that is subject to change.
+        This module must be used with lr scheduler module in `LRScheduler Class
+        <https://www.mindspore.cn/docs/en/master/api_python/mindspore.experimental.html#lrscheduler-class>`_ .
+
+    Args:
+        params (Union[list(Parameter), list(dict)]): list of parameters to optimize or dicts defining
+            parameter groups.
+        lr (Union[int, float, Tensor], optional): learning rate. Default: ``1e-2``.
+        lr_decay (float, optional): learning rate decay. Default: ``0.``.
+        weight_decay (float, optional): weight decay (L2 penalty). Default: ``0.``.
+        initial_accumulator_value (float, optional): the initial accumulator value. Default: ``0.``.
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. Default: ``1e-10``.
+
+    Keyword Args:
+        maximize (bool, optional): maximize the params based on the objective, instead of minimizing.
+            Default: ``False``.
+
+    Inputs:
+        - **gradients** (tuple[Tensor]) - The gradients of `params`.
+
+    Raises:
+        ValueError: If the learning rate is not int, float or Tensor.
+        ValueError: If the learning rate is less than 0.
+        ValueError: If the learning rate decay is less than 0.
+        ValueError: If the `weight_decay` is less than 0.
+        ValueError: If the `initial_accumulator_value` is less than 0.0.
+        ValueError: If the `eps` is less than 0.0.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import nn
+        >>> from mindspore.experimental import optim
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
+        >>> optimizer = optim.Adagrad(net.trainable_params(), lr=0.1)
+        >>> def forward_fn(data, label):
+        ...     logits = net(data)
+        ...     loss = loss_fn(logits, label)
+        ...     return loss, logits
+        >>> grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)
+        >>> def train_step(data, label):
+        ...     (loss, _), grads = grad_fn(data, label)
+        ...     optimizer(grads)
+        ...     return loss
+    """
+
+    def __init__(self, params, lr=1e-2, lr_decay=0.0, weight_decay=0.0, initial_accumulator_value=0.0,
+                 eps=1e-10, *, maximize=False):
+        check_not_less_than_without_equal(lr, "lr", self.cls_name)
+        check_not_less_than(lr_decay, "lr_decay", self.cls_name)
+        check_not_less_than(weight_decay, "weight_decay", self.cls_name)
+        check_not_less_than(initial_accumulator_value, "initial_accumulator_value", self.cls_name)
+        check_not_less_than_without_equal(eps, "eps", self.cls_name)
+
+        defaults = dict(
+            lr=lr,
+            lr_decay=lr_decay,
+            eps=eps,
+            weight_decay=weight_decay,
+            initial_accumulator_value=initial_accumulator_value,
+            maximize=maximize,
+        )
+        super(Adagrad, self).__init__(params, defaults)
+
+        self.accum = self.parameters.clone(prefix="accum", init=initial_accumulator_value)
+        self.op_cast = P.Cast()
+        self.step_t = Parameter(Tensor(0, mstype.int32), "step_t")
+        self.increase_tensor = Tensor(1, mstype.int32)
+        self.assignadd = P.AssignAdd()
+        self.assign = P.Assign()
+
+    @jit
+    def implementation(self, eps, lr, lr_decay, maximize, weight_decay, start_id, end_id, gradients):
+        """Extract the common computing part for acceleration"""
+        opt = P.ApplyAdagradV2(epsilon=eps, update_slots=True)
+        decay_lr = lr / (1 + self.step_t * lr_decay)
+        params = self.parameters[start_id: end_id]
+        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = self._decay_weight(weight_decay, params, grads)
+        accum = self.accum[start_id: end_id]
+        self.hyper_map(F.partial(_adagrad_opt, opt, decay_lr), params, accum, grads)
+        return True
+
+    def construct(self, gradients):
+        for group_id, group in enumerate(self.param_groups):
+
+            lr = self.lrs[group_id]
+            if isinstance(group.get("lr"), float):
+                lr = self.op_cast(group.get("lr"), mstype.float32)
+
+            lr_decay = group["lr_decay"]
+            maximize = group.get("maximize")
+            weight_decay = group["weight_decay"]
+            eps = group["eps"]
+
+            start_id = self.group_start_id[group_id]
+            end_id = self.group_start_id[group_id + 1]
+            self.implementation(eps, lr, lr_decay, maximize, weight_decay, start_id, end_id, gradients)
+
+        self.assignadd(self.step_t, self.increase_tensor)
+
+        return True

mindspore/experimental/optim/adam.py

@@ -20,25 +20,25 @@ from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer
+from mindspore.common.api import jit
 
 _adam_opt = C.MultitypeFuncGraph("adam_opt")
+adam_op = P.Adam(False, False)
 
 
-@_adam_opt.register("Function", "Tensor", "Tensor", "Float", "Float", "Float", "Tensor",
+@_adam_opt.register("Tensor", "Tensor", "Float", "Float", "Float", "Tensor",
                     "Tensor", "Tensor", "Tensor", "Tensor")
-def _run_adam_opt(opt, beta1_power, beta2_power, beta1, beta2, eps, lr, gradient, param, moment1, moment2):
+def _run_adam_opt(beta1_power, beta2_power, beta1, beta2, eps, lr, gradient, param, moment1, moment2):
     """Apply adam optimizer to the weight parameter."""
-    success = True
-    success = F.depend(success, opt(param, moment1, moment2, beta1_power, beta2_power, lr, beta1, beta2, eps, gradient))
-    return success
+    adam_op(param, moment1, moment2, beta1_power, beta2_power, lr, beta1, beta2, eps, gradient)
+    return True
 
 
-@_adam_opt.register("Function", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
-def _run_adam_with_amsgrad_opt(opt, beta1_power, beta2_power, lr, gradient, param, moment1, moment2, vhat):
+@_adam_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
+def _run_adam_with_amsgrad_opt(beta1_power, beta2_power, lr, gradient, param, moment1, moment2, vhat):
     """Apply adam optimizer to the weight parameter with amsgrad."""
-    success = True
-    success = F.depend(success, opt(param, moment1, moment2, vhat, beta1_power, beta2_power, lr, gradient))
-    return success
+    adam_op(param, moment1, moment2, vhat, beta1_power, beta2_power, lr, gradient)
+    return True
 
 
 class Adam(Optimizer):
@@ -80,7 +80,7 @@ class Adam(Optimizer):
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
-        <https://www.mindspore.cn/docs/en/r2.2/api_python/mindspore.nn.html#learningrateschedule-class>`_ .
+        <https://www.mindspore.cn/docs/en/master/api_python/mindspore.nn.html#learningrateschedule-class>`_ .
 
     Args:
         params (Union[list(Parameter), list(dict)]): list of parameters to optimize or dicts defining
@@ -90,7 +90,7 @@ class Adam(Optimizer):
             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: ``0``.
+        weight_decay (float, optional): weight decay (L2 penalty). Default: ``0.``.
         amsgrad (bool, optional): whether to use the AMSGrad algorithm. Default: ``False``.
 
     Keyword Args:
@@ -104,7 +104,7 @@ class Adam(Optimizer):
         ValueError: If the `lr` is not int, float or Tensor.
         ValueError: If the `lr` is less than 0.
         ValueError: If the `eps` is less than 0.0.
-        ValueError: If the `betas` not in the range of 0-1.
+        ValueError: If the `betas` not in the range of [0, 1).
         ValueError: If the `weight_decay` is less than 0.
 
     Supported Platforms:
@@ -115,7 +115,7 @@ class Adam(Optimizer):
         >>> from mindspore import nn
         >>> from mindspore.experimental import optim
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
         >>> optimizer = optim.Adam(net.trainable_params(), lr=0.1)
@@ -129,8 +129,9 @@ class Adam(Optimizer):
         ...     optimizer(grads)
         ...     return loss
     """
+
     def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
-                 weight_decay=0, amsgrad=False, *, maximize=False):
+                 weight_decay=0.0, amsgrad=False, *, maximize=False):
         if lr < 0.0:
             raise ValueError("Invalid learning rate: {}".format(lr))
         if eps < 0.0:
@@ -159,34 +160,34 @@ class Adam(Optimizer):
         self.adam_opt = P.Adam(False, False)
         self.op_cast = P.Cast()
 
+    @jit
+    def implementation(self, beta1, beta2, eps, lr, start_id, end_id, gradients, maximize, weight_decay):
+        """Extract the common computing part for acceleration"""
+        beta1_power = self.op_pow(beta1, self.state_step)
+        beta2_power = self.op_pow(beta2, self.state_step)
+        params = self.parameters[start_id: end_id]
+        grads = tuple([grad if not maximize else F.neg(grad) for grad in gradients[start_id: end_id]])
+        grads = self._decay_weight(weight_decay, params, grads)
+        self.hyper_map(F.partial(_adam_opt, beta1_power, beta2_power, beta1, beta2, eps, lr),
+                       grads, params,
+                       self.exp_avg[start_id: end_id], self.exp_avg_sq[start_id: end_id])
+        return True
+
     def construct(self, gradients):
         self.assignadd(self.state_step, self.increase_tensor)
         for group_id, group in enumerate(self.param_groups):
             start_id = self.group_start_id[group_id]
-            end_id = self.group_start_id[group_id+1]
+            end_id = self.group_start_id[group_id + 1]
 
-            lr = group.get("lr")
+            lr = self.lrs[group_id]
             weight_decay = group.get("weight_decay")
             beta1, beta2 = group.get("betas")
             maximize = group.get("maximize")
             eps = group.get("eps")
 
-            beta1_power = self.op_pow(beta1, self.state_step)
-            beta2_power = self.op_pow(beta2, self.state_step)
-            adam_with_amsgrad_opt = P.ApplyAdamWithAmsgrad(beta1, beta2, eps, False)
-            params = self.parameters[start_id: end_id]
-            grads = gradients[start_id: end_id]
-            grads = grads if not maximize else -grads
-            grads = self._decay_weight(weight_decay, params, grads)
-            if isinstance(lr, float):
+            if isinstance(group.get("lr"), float):
                 lr = self.op_cast(group.get("lr"), mstype.float32)
-            if group.get("amsgrad"):
-                self.hyper_map(F.partial(_adam_opt, adam_with_amsgrad_opt, beta1_power, beta2_power, lr),
-                               grads, params,
-                               self.exp_avg[start_id: end_id], self.exp_avg_sq[start_id: end_id],
-                               self.max_exp_avg_sq[start_id: end_id])
-            else:
-                self.hyper_map(F.partial(_adam_opt, self.adam_opt, beta1_power, beta2_power, beta1, beta2, eps, lr),
-                               grads, params,
-                               self.exp_avg[start_id: end_id], self.exp_avg_sq[start_id: end_id])
+
+            self.implementation(beta1, beta2, eps, lr, start_id, end_id, gradients, maximize, weight_decay)
+
         return True