mindspore 2.2.11__cp39-cp39-win_amd64.whl → 2.3.0__cp39-cp39-win_amd64.whl

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

mindspore/experimental/optim/adamax.py

@@ -0,0 +1,170 @@
+# 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.
+# ============================================================================
+"""adamax"""
+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 import _checkparam as validator
+from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_than, check_not_less_than_without_equal
+from mindspore import ops
+from mindspore import jit
+
+_adamax_opt = C.MultitypeFuncGraph("adamax_opt")
+
+
+@_adamax_opt.register("Number", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
+def _tensor_run_opt(beta1, beta2, eps, clr, param, grad, exp_avg, exp_inf):
+    """Apply adamax optimizer to the weight parameter."""
+    F.assign(exp_avg, exp_avg * beta1 + grad * (1-beta1))
+    norm_buf = ops.cat([ops.unsqueeze(exp_inf * beta2, 0), ops.unsqueeze(grad.abs().add(eps), 0)], 0)
+    F.assign(exp_inf, ops.amax(norm_buf, 0))
+
+    F.assign(param, param - clr * exp_avg / exp_inf)
+    return True
+
+
+class Adamax(Optimizer):
+    r"""
+    Implements Adamax algorithm (a variant of Adam based on infinity norm).
+
+    .. math::
+       \begin{aligned}
+            &\rule{110mm}{0.4pt}                                                                 \\
+            &\textbf{input}      : \gamma \text{ (lr)}, \beta_1, \beta_2
+                \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)},
+                \: \lambda \text{ (weight decay)},                                                \\
+            &\hspace{13mm}    \epsilon \text{ (epsilon)}                                          \\
+            &\textbf{initialize} :  m_0 \leftarrow 0 \text{ ( first moment)},
+                u_0 \leftarrow 0 \text{ ( infinity norm)}                                 \\[-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}m_t      \leftarrow   \beta_1 m_{t-1} + (1 - \beta_1) g_t               \\
+            &\hspace{5mm}u_t      \leftarrow   \mathrm{max}(\beta_2 u_{t-1}, |g_{t}|+\epsilon)   \\
+            &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \frac{\gamma m_t}{(1-\beta^t_1) u_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: ``2e-3``.
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square. 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.``.
+
+    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 0.0.
+        ValueError: If the `weight_decay` is less than 0.
+        ValueError: If elements of the `betas` not in the range of [0,1).
+
+    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.Adamax(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=2e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.0, *, maximize=False):
+        check_not_less_than_without_equal(lr, "lr", self.cls_name)
+        check_not_less_than(weight_decay, "weight_decay", self.cls_name)
+        check_not_less_than_without_equal(eps, "eps", self.cls_name)
+        validator.check_float_range(betas[0], 0., 1., validator.INC_LEFT, "betas[0]", self.cls_name)
+        validator.check_float_range(betas[1], 0., 1., validator.INC_LEFT, "betas[1]", self.cls_name)
+
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            maximize=maximize,
+        )
+        super(Adamax, self).__init__(params, defaults)
+
+        self.step_t = Parameter(Tensor(0, mstype.int32), "step_t")
+        self.exp_avg = self.parameters.clone(prefix="exp_avg", init='zeros')
+        self.exp_inf = self.parameters.clone(prefix="exp_inf", init='zeros')
+        self.increase_tensor = Tensor(1, mstype.int32)
+        self.assignadd = P.AssignAdd()
+        self.op_cast = P.Cast()
+
+    @jit
+    def implementation(self, group_id, lr, gradients, maximize, weight_decay, beta1, beta2, eps):
+        """Extract the common computing part for acceleration"""
+        start_id = self.group_start_id[group_id]
+        end_id = self.group_start_id[group_id + 1]
+        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)
+        exp_avg = self.exp_avg[start_id: end_id]
+        exp_inf = self.exp_inf[start_id: end_id]
+        bias_correction = 1 - beta1 ** self.step_t
+        clr = lr / bias_correction
+        self.hyper_map(F.partial(_adamax_opt, beta1, beta2, eps, clr),
+                       params, grads, exp_avg, exp_inf)
+        return True
+
+    def construct(self, gradients):
+        self.assignadd(self.step_t, self.increase_tensor)
+
+        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")
+            beta1, beta2 = group["betas"]
+
+            self.implementation(group_id, lr, gradients, maximize, group["weight_decay"], beta1, beta2, group["eps"])
+        return True

mindspore/experimental/optim/adamw.py

@@ -21,6 +21,7 @@ from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
 from mindspore.experimental.optim.optimizer import Optimizer
 from mindspore import ops
+from mindspore import jit
 
 _adamw_opt = C.MultitypeFuncGraph("adamw_opt")
 
@@ -28,21 +29,28 @@ op_mul = P.Mul()
 op_pow = P.Pow()
 op_sqrt = P.Sqrt()
 op_maximum = P.Maximum()
+hyper_map = C.HyperMap()
 
 
-@_adamw_opt.register("Float", "Tensor", "Bool", "Float", "Tensor", "Float", "Float", "Tensor", "Tensor",
+@jit
+def prepare_func(lr, weight_decay, state_step, beta1, beta2):
+    weight_decay_new = 1 - lr * weight_decay
+    bias_correction1 = 1 - op_pow(beta1, state_step)
+    bias_correction2 = 1 - op_pow(beta2, state_step)
+    step_size = lr / bias_correction1
+    bias_correction2_sqrt = op_sqrt(bias_correction2)
+    return weight_decay_new, step_size, bias_correction2_sqrt
+
+
+@_adamw_opt.register("Tensor", "Tensor", "Bool", "Float", "Tensor", "Float", "Float", "Tensor", "Tensor",
                      "Tensor", "Tensor", "Tensor")
-def _run_adamw_opt(weight_decay, lr, amsgrad, eps, state_step, beta1, beta2, param, grad,
+def _run_adamw_opt(weight_decay_new, step_size, amsgrad, eps, bias_correction2_sqrt, beta1, beta2, param, grad,
                    exp_avg, exp_avg_sq, max_exp_avg_sq):
     """Apply adamw optimizer to the weight parameter."""
     success = True
-    next_param = op_mul(param, 1 - lr * weight_decay)
+    next_param = op_mul(param, weight_decay_new)
     F.assign(exp_avg, op_mul(exp_avg, beta1) + op_mul(grad, 1 - beta1))
     F.assign(exp_avg_sq, ops.addcmul(op_mul(exp_avg_sq, beta2), grad, grad, 1 - beta2))
-    bias_correction1 = 1 - op_pow(beta1, state_step)
-    bias_correction2 = 1 - op_pow(beta2, state_step)
-    step_size = lr / bias_correction1
-    bias_correction2_sqrt = op_sqrt(bias_correction2)
 
     if amsgrad:
         next_max_exp_avg = op_maximum(max_exp_avg_sq, exp_avg_sq)
@@ -93,7 +101,7 @@ class AdamW(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.experimental.html#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
@@ -103,7 +111,7 @@ class AdamW(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:
@@ -117,7 +125,7 @@ class AdamW(Optimizer):
         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 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:
@@ -128,7 +136,7 @@ class AdamW(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.AdamW(net.trainable_params(), lr=0.1)
@@ -142,6 +150,7 @@ class AdamW(Optimizer):
         ...     optimizer(grads)
         ...     return loss
     """
+
     def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
                  weight_decay=1e-2, amsgrad=False, *, maximize=False):
         if lr < 0.0:
@@ -166,6 +175,18 @@ class AdamW(Optimizer):
         self.state_step = Parameter(Tensor(0, mstype.int32), "state_step")
         self.increase_tensor = Tensor(1, mstype.int32)
         self.assignadd = P.AssignAdd()
+        self.op_cast = P.Cast()
+
+    @jit
+    def implementation(self, lr, weight_decay, beta1, beta2, amsgrad, eps, grads, start_id, end_id):
+        """Extract the common computing part for acceleration"""
+        weight_decay_new, step_size, bias_correction2_sqrt = prepare_func(lr, weight_decay,
+                                                                          self.state_step, beta1, beta2)
+        self.hyper_map(F.partial(_adamw_opt, weight_decay_new, step_size, amsgrad,
+                                 eps, bias_correction2_sqrt, beta1, beta2),
+                       self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
+                       self.exp_avg_sq[start_id: end_id], self.max_exp_avg_sq[start_id: end_id])
+        return True
 
     def construct(self, gradients):
         self.assignadd(self.state_step, self.increase_tensor)
@@ -173,9 +194,12 @@ class AdamW(Optimizer):
             beta1, beta2 = group['betas']
             start_id = self.group_start_id[group_id]
             end_id = self.group_start_id[group_id + 1]
-            grads = gradients[start_id: end_id] if not group.get("maximize") else -gradients[start_id: end_id]
-            self.hyper_map(F.partial(_adamw_opt, group.get("weight_decay"), group.get("lr"), group.get("amsgrad"),
-                                     group.get("eps"), self.state_step, beta1, beta2),
-                           self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
-                           self.exp_avg_sq[start_id: end_id], self.max_exp_avg_sq[start_id: end_id])
+            lr = self.lrs[group_id]
+            if isinstance(group.get("lr"), float):
+                lr = self.op_cast(group.get("lr"), mstype.float32)
+            grads = tuple([grad if not group.get("maximize") else F.neg(grad) for grad in gradients[start_id: end_id]])
+
+            self.implementation(lr, group.get("weight_decay"), beta1, beta2, group.get("amsgrad"), group.get("eps"),
+                                grads, start_id, end_id)
+
         return True

mindspore/experimental/optim/asgd.py

@@ -0,0 +1,153 @@
+# Copyright 2021-2022 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.
+# ============================================================================
+"""asgd"""
+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.common.api import jit
+
+_asgd_opt = C.MultitypeFuncGraph("asgd_opt")
+
+op_cast = P.Cast()
+op_pow = P.Pow()
+op_maximum = P.Maximum()
+op_assign = P.Assign()
+op_assignadd = P.AssignAdd()
+
+
+@_asgd_opt.register("Number", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor",
+                    "Tensor", "Tensor", "Tensor")
+def _run_asgd_opt(lambd, alpha, t0, step, lr, param, grad, eta, mu, ax):
+    """Apply asgd optimizer to the weight parameter using dynamic learning rate."""
+    if step == 1:
+        op_assign(eta, lr)
+    next_param = op_cast(param * (1. - lambd * eta) - eta * grad, param.dtype)
+    F.assign(param, next_param)
+
+    if mu != 1:
+        op_assignadd(ax, op_cast((next_param - ax) * mu, ax.dtype))
+    else:
+        op_assign(ax, next_param)
+
+    op_assign(eta, lr / (op_pow((1. + lambd * lr * step), alpha)))
+    op_assign(mu, 1. / op_maximum(1., step - t0))
+    return True
+
+
+class ASGD(Optimizer):
+    r"""
+    Implements Averaged Stochastic Gradient Descent algorithm.
+
+    .. 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``.
+        lambd (float, optional): decay term. Default: ``1e-4``.
+        alpha (float, optional): power for eta update. Default: ``0.75``.
+        t0 (float, optional): point at which to start averaging. Default: ``1e6``.
+        weight_decay (float, optional): weight decay (L2 penalty). Default: ``0.``.
+        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 `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.ASGD(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, lambd=1e-4, alpha=0.75, t0=1e6, weight_decay=0.0, maximize=False):
+        check_not_less_than_without_equal(lr, "lr", self.cls_name)
+        check_not_less_than(weight_decay, "weight_decay", self.cls_name)
+        if not isinstance(lambd, float):
+            raise TypeError(f"For 'ASGD', the type of lambd must be float, but got {type(lambd)}.")
+        if not isinstance(t0, float):
+            raise TypeError(f"For 'ASGD', the type of t0 must be float, but got {type(t0)}.")
+        defaults = dict(
+            lr=lr,
+            lambd=lambd,
+            alpha=alpha,
+            t0=t0,
+            weight_decay=weight_decay,
+            maximize=maximize,
+        )
+        super(ASGD, self).__init__(params, defaults)
+        self.mu = [Parameter(Tensor(1.), "mu_" + param.name) for param in self.parameters]
+        self.eta = [Parameter(Tensor(0.), "eta_" + param.name) for param in self.parameters]
+        self.ax = self.parameters.clone(prefix="ax", init='zeros')
+        self.step_t = Parameter(Tensor(0, mstype.int32), "step_t")
+        self.increase_tensor = Tensor(1, mstype.int32)
+        self.assignadd = P.AssignAdd()
+        self.op_cast = P.Cast()
+
+    @jit
+    def implementation(self, lambd, alpha, t0, lr, group_id, maximize, gradients, weight_decay):
+        """Extract the common computing part for acceleration"""
+        start_id = self.group_start_id[group_id]
+        end_id = self.group_start_id[group_id + 1]
+        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)
+
+        ax = self.ax[start_id: end_id]
+        eta = self.eta[start_id: end_id]
+        mu = self.mu[start_id: end_id]
+        self.hyper_map(F.partial(_asgd_opt, lambd, alpha, t0, self.step_t, lr),
+                       params, grads, eta, mu, ax)
+        return True
+
+    def construct(self, gradients):
+        self.assignadd(self.step_t, self.increase_tensor)
+        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")
+
+            self.implementation(group["lambd"], group["alpha"], group["t0"], lr, group_id, maximize, gradients,
+                                group["weight_decay"])
+        return True