mindspore 2.6.0rc1__cp311-cp311-win_amd64.whl → 2.7.0__cp311-cp311-win_amd64.whl

mindspore/experimental/optim/adamw.py

@@ -1,21 +1,10 @@
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/adamw.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adamw"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 import mindspore.common.dtype as mstype
@@ -25,14 +14,14 @@ from mindspore.ops import auto_generate as gen
 from mindspore import ops
 from mindspore import jit
 
-_adamw_opt = C.MultitypeFuncGraph("adamw_opt")
-_speed_adamw_opt = C.MultitypeFuncGraph("speed_adamw_opt")
+_adamw_opt = ops.MultitypeFuncGraph("adamw_opt")
+_speed_adamw_opt = ops.MultitypeFuncGraph("speed_adamw_opt")
 
-op_mul = P.Mul()
-op_pow = P.Pow()
-op_sqrt = P.Sqrt()
-op_maximum = P.Maximum()
-hyper_map = C.HyperMap()
+op_mul = ops.Mul()
+op_pow = ops.Pow()
+op_sqrt = ops.Sqrt()
+op_maximum = ops.Maximum()
+hyper_map = ops.HyperMap()
 
 
 @_speed_adamw_opt.register("Function", "Float", "Float", "Tensor", "Float", "Float", "Bool", "Bool", "Tensor", "Tensor",
@@ -76,18 +65,18 @@ def _run_adamw_opt(weight_decay_new, step_size, amsgrad, eps, bias_correction2_s
     """Apply adamw optimizer to the weight parameter."""
     success = True
     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))
+    ops.assign(exp_avg, op_mul(exp_avg, beta1) + op_mul(grad, 1 - beta1))
+    ops.assign(exp_avg_sq, ops.addcmul(op_mul(exp_avg_sq, beta2), grad, grad, 1 - beta2))
 
     if amsgrad:
         next_max_exp_avg = op_maximum(max_exp_avg_sq, exp_avg_sq)
         denom = op_sqrt(next_max_exp_avg) / bias_correction2_sqrt + eps
-        F.assign(max_exp_avg_sq, next_max_exp_avg)
+        ops.assign(max_exp_avg_sq, next_max_exp_avg)
     else:
         denom = op_sqrt(exp_avg_sq) / bias_correction2_sqrt + eps
 
     return_param = next_param - op_mul(exp_avg / denom, step_size)
-    F.assign(param, return_param)
+    ops.assign(param, return_param)
     return success
 
 
@@ -129,6 +118,10 @@ class AdamW(Optimizer):
             &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
+    More details of the AdamW algorithm can be found in the paper `Decoupled Weight Decay Regularization
+    <https://arxiv.org/abs/1711.05101>`_ and `On the Convergence of Adam and Beyond
+    <https://openreview.net/forum?id=ryQu7f-RZ>`_.
+
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
@@ -205,16 +198,16 @@ class AdamW(Optimizer):
         self.max_exp_avg_sq = self.parameters.clone(prefix="max_exp_avg_sq", init='zeros')
         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()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.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.hyper_map(ops.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
@@ -228,7 +221,8 @@ class AdamW(Optimizer):
             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]])
+            grads = tuple([grad if not group.get("maximize") else ops.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)
@@ -265,7 +259,7 @@ class SpeedAdamW(Optimizer):
         self.exp_avg_sq = self.parameters.clone(prefix="exp_avg_sq", init='zeros')
         self.state_step = Parameter(Tensor([0], mstype.float32), "state_step")
         self.increase_tensor = Tensor(1, mstype.float32)
-        self.assignadd = P.AssignAdd()
+        self.assignadd = ops.AssignAdd()
         self.adamw_opt = gen.ApplyAdamW()
 
     def construct(self, gradients):
@@ -285,9 +279,9 @@ class SpeedAdamW(Optimizer):
             if group.get("amsgrad"):
                 raise ValueError("For SpeedAdamW, the value of amsgrad can only be False.")
 
-            self.hyper_map(F.partial(_speed_adamw_opt, self.adamw_opt, beta1, beta2, lr,
-                                     group.get("eps"), group.get("weight_decay"),
-                                     group.get("amsgrad"), maximize, bias_correction1, bias_correction2),
+            self.hyper_map(ops.partial(_speed_adamw_opt, self.adamw_opt, beta1, beta2, lr,
+                                       group.get("eps"), group.get("weight_decay"),
+                                       group.get("amsgrad"), maximize, bias_correction1, bias_correction2),
                            self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
                            self.exp_avg_sq[start_id: end_id])
 

mindspore/experimental/optim/asgd.py

@@ -1,33 +1,23 @@
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/asgd.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """asgd"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 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")
+_asgd_opt = ops.MultitypeFuncGraph("asgd_opt")
 
-op_cast = P.Cast()
-op_pow = P.Pow()
-op_maximum = P.Maximum()
-op_assign = P.Assign()
-op_assignadd = P.AssignAdd()
+op_cast = ops.Cast()
+op_pow = ops.Pow()
+op_maximum = ops.Maximum()
+op_assign = ops.Assign()
+op_assignadd = ops.AssignAdd()
 
 
 @_asgd_opt.register("Number", "Number", "Number", "Tensor", "Tensor", "Tensor", "Tensor",
@@ -37,7 +27,7 @@ def _run_asgd_opt(lambd, alpha, t0, step, lr, param, grad, eta, mu, ax):
     if step == 1:
         op_assign(eta, lr)
     next_param = op_cast(param * (1. - lambd * eta) - eta * grad, param.dtype)
-    F.assign(param, next_param)
+    ops.assign(param, next_param)
 
     if mu != 1:
         op_assignadd(ax, op_cast((next_param - ax) * mu, ax.dtype))
@@ -121,8 +111,8 @@ class ASGD(Optimizer):
         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()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.Cast()
 
     @jit(backend="ms_backend")
     def implementation(self, lambd, alpha, t0, lr, group_id, maximize, gradients, weight_decay):
@@ -130,13 +120,13 @@ class ASGD(Optimizer):
         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 = tuple([grad if not maximize else ops.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),
+        self.hyper_map(ops.partial(_asgd_opt, lambd, alpha, t0, self.step_t, lr),
                        params, grads, eta, mu, ax)
         return True
 

mindspore/experimental/optim/lr_scheduler.py

@@ -1,16 +1,6 @@
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/lr_scheduler.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """LRScheduler."""
 from collections import Counter
@@ -20,8 +10,6 @@ from mindspore import ops, Tensor, Parameter
 from mindspore.experimental.optim.optimizer import Optimizer
 from mindspore.common.api import jit_class
 import mindspore.common.dtype as mstype
-from mindspore.ops import functional as F
-from mindspore.ops import operations as P
 from mindspore import _checkparam as Validator
 
 __all__ = ['StepLR', 'LinearLR', 'LRScheduler', 'ExponentialLR', 'PolynomialLR',
@@ -143,9 +131,12 @@ class LRScheduler:
 
 @jit_class
 class StepLR(LRScheduler):
-    """Decays the learning rate of each parameter group by gamma every
-    step_size epochs. Notice that such decay can happen simultaneously with
-    other changes to the learning rate from outside this scheduler.
+    """
+    During training, when calling `StepLR.step()` , if the current epoch number is an integer multiple of `step_size` ,
+    the learning rate will be decayed by multiplying it with `gamma` . The adjustment of the learning rate and
+    the parameter update of the optimizer are synergistically performed. The optimizer executes parameter optimization
+    operations based on the currently adjusted learning rate. The learning rate decay of StepLR may occur simultaneously
+    with external changes to the learning rate.
 
     .. warning::
         This is an experimental lr scheduler module that is subject to change.
@@ -431,8 +422,8 @@ class PolynomialLR(LRScheduler):
             raise TypeError(f"For 'PolynomialLR', the type of total_iters must be int, but got {type(total_iters)}.")
         self.total_iters = total_iters
         self.power = power
-        self.min = P.Minimum()
-        self.cast = P.Cast()
+        self.min = ops.Minimum()
+        self.cast = ops.Cast()
         super(PolynomialLR, self).__init__(optimizer, last_epoch)
 
     def get_lr(self):
@@ -804,7 +795,7 @@ class SequentialLR:
 
 @jit_class
 class ReduceLROnPlateau:
-    """
+    r"""
     Reduce learning rate when a metric has stopped improving.
     Models often benefit from reducing the learning rate by a factor
     of 2-10 once learning stagnates. The scheduler reads the metrics `metrics` during execution
@@ -886,7 +877,7 @@ class ReduceLROnPlateau:
         [Tensor(shape=[], dtype=Float32, value= 0.001)]
         [Tensor(shape=[], dtype=Float32, value= 0.001)]
         [Tensor(shape=[], dtype=Float32, value= 0.0001)]
-        """
+    """
 
     def __init__(self, optimizer, mode='min', factor=0.1, patience=10,
                  threshold=1e-4, threshold_mode='rel', cooldown=0,
@@ -915,8 +906,8 @@ class ReduceLROnPlateau:
         self.cooldown_counter = 0
         self.eps = eps
         self.mode_worse = None
-        self.assign = P.Assign()
-        self.cast = P.Cast()
+        self.assign = ops.Assign()
+        self.cast = ops.Cast()
         self.last_epoch = Parameter(Tensor(0, dtype=mstype.int32),
                                     name='last_epoch_' + self.__class__.__name__)
 
@@ -1079,17 +1070,8 @@ class CyclicLR(LRScheduler):
         [Tensor(shape=[], dtype=Float32, value= 0.010225)]
     """
 
-    def __init__(self,
-                 optimizer,
-                 base_lr,
-                 max_lr,
-                 step_size_up=2000,
-                 step_size_down=None,
-                 mode='triangular',
-                 gamma=1.,
-                 scale_fn=None,
-                 scale_mode='cycle',
-                 last_epoch=-1):
+    def __init__(self, optimizer, base_lr, max_lr, step_size_up=2000, step_size_down=None, mode='triangular',
+                 gamma=1.0, scale_fn=None, scale_mode='cycle', last_epoch=-1):
 
         base_lrs = self._preprocess_input_param(optimizer, base_lr, 'base_lr')
 
@@ -1117,7 +1099,7 @@ class CyclicLR(LRScheduler):
         self._scale_fn_custom = scale_fn
         self.scale_mode = scale_mode
         self._init_scale_fn()
-        self.floor = P.Floor()
+        self.floor = ops.Floor()
 
         super(CyclicLR, self).__init__(optimizer, last_epoch)
         self.base_lrs = [Tensor(lr) for lr in base_lrs]
@@ -1252,12 +1234,12 @@ class CosineAnnealingWarmRestarts(LRScheduler):
         self.zero_tensor = Tensor(0, mstype.int32)
 
         self.math_pi = math.pi
-        self.cos = P.Cos()
-        self.cast = P.Cast()
-        self.log = P.Log()
-        self.cast = P.Cast()
-        self.assign = P.Assign()
-        self.floor = P.Floor()
+        self.cos = ops.Cos()
+        self.cast = ops.Cast()
+        self.log = ops.Log()
+        self.cast = ops.Cast()
+        self.assign = ops.Assign()
+        self.floor = ops.Floor()
         self._last_lr = [group["lr"] for group in optimizer.param_groups]
         super(CosineAnnealingWarmRestarts, self).__init__(optimizer, last_epoch)
 
@@ -1306,7 +1288,7 @@ class CosineAnnealingWarmRestarts(LRScheduler):
 
         for i, data in enumerate(zip(self.optimizer.param_groups, self.get_lr())):
             _, lr = data
-            F.assign(self.optimizer.param_groups[i]["lr"], lr)
+            ops.assign(self.optimizer.param_groups[i]["lr"], lr)
 
 
 @jit_class
@@ -1371,8 +1353,8 @@ class CosineAnnealingLR(LRScheduler):
         self.T_max = T_max
         self.eta_min = eta_min
         self.math_pi = math.pi
-        self.cos = P.Cos()
-        self.cast = P.Cast()
+        self.cos = ops.Cos()
+        self.cast = ops.Cast()
         super(CosineAnnealingLR, self).__init__(optimizer, last_epoch)
 
     def get_lr(self):

mindspore/experimental/optim/nadam.py

@@ -1,30 +1,20 @@
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/nadam.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """nadam"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 from mindspore.common import Parameter, Tensor
 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 jit
 
-_nadam_opt = C.MultitypeFuncGraph("nadam_opt")
+_nadam_opt = ops.MultitypeFuncGraph("nadam_opt")
 
-op_sqrt = P.Sqrt()
+op_sqrt = ops.Sqrt()
 
 
 @_nadam_opt.register("Number", "Number", "Number", "Number", "Tensor", "Tensor", "Tensor",
@@ -34,15 +24,15 @@ def _tensor_run_opt(beta1, beta2, momentum_decay, eps, step_t, lr, param, grad,
     bias_correction2 = 1 - beta2 ** step_t
     mu = beta1 * (1. - 0.5 * (0.96 ** (step_t * momentum_decay)))
     mu_next = beta1 * (1. - 0.5 * (0.96 ** ((step_t + 1) * momentum_decay)))
-    F.assign(mu_product, mu_product * mu)
-    F.assign(exp_avg, exp_avg * beta1 + grad * (1 - beta1))
-    F.assign(exp_avg_sq, exp_avg_sq * beta2 + grad * grad * (1 - beta2))
+    ops.assign(mu_product, mu_product * mu)
+    ops.assign(exp_avg, exp_avg * beta1 + grad * (1 - beta1))
+    ops.assign(exp_avg_sq, exp_avg_sq * beta2 + grad * grad * (1 - beta2))
 
     denom = op_sqrt(exp_avg_sq / bias_correction2) + eps
 
     mu_product_next = mu_product * mu_next
-    F.assign(param, param - lr * (1. - mu) / (1. - mu_product) * grad / denom)
-    F.assign(param, param - (lr * mu_next) / (1. - mu_product_next) * exp_avg / denom)
+    ops.assign(param, param - lr * (1. - mu) / (1. - mu_product) * grad / denom)
+    ops.assign(param, param - (lr * mu_next) / (1. - mu_product_next) * exp_avg / denom)
 
     return True
 
@@ -122,8 +112,8 @@ class NAdam(Optimizer):
         self.mu_product = [Parameter(Tensor(1.), "mu_product_" + param.name) for param in self.parameters]
 
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, lr, beta1, beta2, weight_decay, momentum_decay, eps, start_id, end_id, gradients):
@@ -135,7 +125,7 @@ class NAdam(Optimizer):
         exp_avg_sq = self.exp_avg_sq[start_id: end_id]
         mu_product = self.mu_product[start_id: end_id]
 
-        self.hyper_map(F.partial(_nadam_opt, beta1, beta2, momentum_decay, eps, self.step_t, lr),
+        self.hyper_map(ops.partial(_nadam_opt, beta1, beta2, momentum_decay, eps, self.step_t, lr),
                        params, grads, exp_avg, exp_avg_sq, mu_product)
         return True
 

mindspore/experimental/optim/optimizer.py

@@ -1,22 +1,12 @@
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/optimizer.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """optimizer"""
 from __future__ import absolute_import
 from collections import defaultdict
 from typing import Iterable
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 
 from mindspore.nn.cell import Cell
 from mindspore.common.parameter import Parameter, ParameterTuple
@@ -98,7 +88,7 @@ class Optimizer(Cell):
         self.param_groups = []
         self.parameters = []
         self.lrs = []
-        self.map_ = C.Map()
+        self.map_ = ops.Map()
         self.group_start_id = [0]
         if not isinstance(param_groups[0], dict):
             param_groups = [{'params': param_groups}]
@@ -106,7 +96,7 @@ class Optimizer(Cell):
         for param_group in param_groups:
             self.add_param_group(param_group)
         self.parameters = ParameterTuple(self.parameters)
-        self.hyper_map = C.HyperMap()
+        self.hyper_map = ops.HyperMap()
         self.enable_tuple_broaden = True
 
     def __repr__(self):
@@ -167,7 +157,7 @@ class Optimizer(Cell):
         """Apply weight decay."""
         if weight_decay != 0.:
             weight_decay = Tensor(weight_decay, mstype.float32)
-            gradients = self.map_(F.partial(_apply_decay, weight_decay), params, gradients)
+            gradients = self.map_(ops.partial(_apply_decay, weight_decay), params, gradients)
         return gradients
 
     def _preprocess_param_group(self, param_group):
@@ -228,18 +218,18 @@ class Optimizer(Cell):
     def construct(self, *hyper_params):
         raise NotImplementedError
 
-op_add = P.AddN()
-op_gather = P.Gather()
-op_mul = P.Mul()
+op_add = ops.AddN()
+op_gather = ops.Gather()
+op_mul = ops.Mul()
 
-_apply_decay = C.MultitypeFuncGraph("apply_decay")
+_apply_decay = ops.MultitypeFuncGraph("apply_decay")
 
 
 @_apply_decay.register("Tensor", "Tensor", "RowTensor")
 def _tensor_apply_decay_with_sparse(weight_decay, weight, gradient):
     """Get grad with weight_decay."""
     indices = gradient.indices
-    values = op_add((op_gather(weight, indices, 0) * F.cast(weight_decay, F.dtype(weight)), gradient.values))
+    values = op_add((op_gather(weight, indices, 0) * ops.cast(weight_decay, ops.dtype(weight)), gradient.values))
     shape = gradient.dense_shape
     return RowTensorInner(indices, values, shape)
 
@@ -247,7 +237,7 @@ def _tensor_apply_decay_with_sparse(weight_decay, weight, gradient):
 @_apply_decay.register("Tensor", "Tensor", "Tensor")
 def _tensor_apply_decay(weight_decay, weight, gradient):
     """Get grad with weight_decay."""
-    return op_add((op_mul(weight, F.cast(weight_decay, F.dtype(weight))), gradient))
+    return op_add((op_mul(weight, ops.cast(weight_decay, ops.dtype(weight))), gradient))
 
 
 def check_not_less_than(arg_value, arg_name, prim, value=0.0):

mindspore/experimental/optim/radam.py

@@ -1,32 +1,22 @@
-# 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.
+# The code implementation refers to the following files from pytorch:
+# - https://github.com/pytorch/pytorch/blob/v1.13.0/torch/optim/radam.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """radam"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore import ops
 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 jit
 
-_radam_opt = C.MultitypeFuncGraph("radam_opt")
+_radam_opt = ops.MultitypeFuncGraph("radam_opt")
 
-op_pow = P.Pow()
-op_sqrt = P.Sqrt()
-op_cast = P.Cast()
+op_pow = ops.Pow()
+op_sqrt = ops.Sqrt()
+op_cast = ops.Cast()
 
 
 @_radam_opt.register("Number", "Number", "Number", "Tensor", "Number", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor",
@@ -35,17 +25,17 @@ def _tensor_run_opt(beta1, beta2, eps, lr, rho_inf, rho_t, bias_correction1, bia
                     exp_avg_sq):
     """Apply radam optimizer to the weight parameter."""
 
-    F.assign(exp_avg, exp_avg * beta1 + grad * (1 - beta1))
-    F.assign(exp_avg_sq, exp_avg_sq * beta2 + grad * grad * (1 - beta2))
+    ops.assign(exp_avg, exp_avg * beta1 + grad * (1 - beta1))
+    ops.assign(exp_avg_sq, exp_avg_sq * beta2 + grad * grad * (1 - beta2))
     bias_corrected_exp_avg = exp_avg / bias_correction1
 
     if rho_t > 5.0:
         rect = op_sqrt((rho_t - 4) * (rho_t - 2) * rho_inf / ((rho_inf - 4) * (rho_inf - 2) * rho_t))
         exp_avg_sq_sqrt = op_sqrt(exp_avg_sq) + eps
         adaptive_lr = op_sqrt(bias_correction2) / exp_avg_sq_sqrt
-        F.assign(param, param - bias_corrected_exp_avg * lr * adaptive_lr * rect)
+        ops.assign(param, param - bias_corrected_exp_avg * lr * adaptive_lr * rect)
     else:
-        F.assign(param, param - bias_corrected_exp_avg * lr)
+        ops.assign(param, param - bias_corrected_exp_avg * lr)
 
     return True
 
@@ -89,6 +79,9 @@ class RAdam(Optimizer):
             &\rule{180mm}{0.4pt}
         \end{align*}
 
+    For more details about RAdam algorithm, please refer to `On the Variance of the Adaptive Learning Rate and Beyond
+    <https://arxiv.org/abs/1908.03265>`_.
+
     .. warning::
         This is an experimental optimizer API that is subject to change.
         This module must be used with lr scheduler module in `LRScheduler Class
@@ -155,7 +148,7 @@ class RAdam(Optimizer):
         self.exp_avg = self.parameters.clone(prefix="exp_avg", init='zeros')
         self.exp_avg_sq = self.parameters.clone(prefix="exp_avg_sq", init='zeros')
         self.increase_tensor = Tensor(1, mstype.int32)
-        self.assignadd = P.AssignAdd()
+        self.assignadd = ops.AssignAdd()
 
     @jit(backend="ms_backend")
     def implementation(self, lr, beta1, beta2, weight_decay, eps, start_id, end_id, gradients):
@@ -175,7 +168,8 @@ class RAdam(Optimizer):
 
         rho_t = rho_inf - right
 
-        self.hyper_map(F.partial(_radam_opt, beta1, beta2, eps, lr, rho_inf, rho_t, bias_correction1, bias_correction2),
+        self.hyper_map(ops.partial(_radam_opt, beta1, beta2, eps, lr, rho_inf,
+                                   rho_t, bias_correction1, bias_correction2),
                        params, grads, exp_avg, exp_avg_sq)
         return True