mindspore 2.6.0__cp39-cp39-win_amd64.whl → 2.7.0rc1__cp39-cp39-win_amd64.whl

mindspore/experimental/optim/adam.py

@@ -1,29 +1,19 @@
-# 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/adam.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """adam"""
 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.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 = ops.MultitypeFuncGraph("adam_opt")
+adam_op = ops.Adam(False, False)
 
 
 @_adam_opt.register("Tensor", "Tensor", "Float", "Float", "Float", "Tensor",
@@ -81,6 +71,9 @@ class Adam(Optimizer):
             &\rule{180mm}{0.4pt}                                                          \\[-1.ex]
        \end{aligned}
 
+    For more details about Adam algorithm, please refer to `Adam: A Method for Stochastic Optimization
+    <https://arxiv.org/abs/1412.6980>`_.
+
     .. warning::
         The implementation formula of this optimizer interface is not completely consistent with that in the paper.
         If you want to use an interface that is completely consistent, it is recommended to use
@@ -160,12 +153,12 @@ class Adam(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_add = P.AddN()
-        self.op_mul = P.Mul()
-        self.op_pow = P.Pow()
-        self.adam_opt = P.Adam(False, False)
-        self.op_cast = P.Cast()
+        self.assignadd = ops.AssignAdd()
+        self.op_add = ops.AddN()
+        self.op_mul = ops.Mul()
+        self.op_pow = ops.Pow()
+        self.adam_opt = ops.Adam(False, False)
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, beta1, beta2, eps, lr, start_id, end_id, gradients, maximize, weight_decay):
@@ -173,9 +166,9 @@ class Adam(Optimizer):
         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 = 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)
-        self.hyper_map(F.partial(_adam_opt, beta1_power, beta2_power, beta1, beta2, eps, lr),
+        self.hyper_map(ops.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

mindspore/experimental/optim/adamax.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/adamax.py
+# Additional modifications are made by Huawei Technologies Co., Ltd in 2023.
 # ============================================================================
 """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
@@ -23,17 +12,17 @@ from mindspore.experimental.optim.optimizer import Optimizer, check_not_less_tha
 from mindspore import ops
 from mindspore import jit
 
-_adamax_opt = C.MultitypeFuncGraph("adamax_opt")
+_adamax_opt = ops.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))
+    ops.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))
+    ops.assign(exp_inf, ops.amax(norm_buf, 0))
 
-    F.assign(param, param - clr * exp_avg / exp_inf)
+    ops.assign(param, param - clr * exp_avg / exp_inf)
     return True
 
 
@@ -68,6 +57,9 @@ class Adamax(Optimizer):
         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>`_ .
 
+    For more details about Adamax algorithm, please refer to `Adam: A Method for Stochastic Optimization
+    <https://arxiv.org/abs/1412.6980>`_.
+
     Args:
         params (Union[list(Parameter), list(dict)]): list of parameters to optimize or dicts defining
             parameter groups.
@@ -135,8 +127,8 @@ class Adamax(Optimizer):
         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()
+        self.assignadd = ops.AssignAdd()
+        self.op_cast = ops.Cast()
 
     @jit
     def implementation(self, group_id, lr, gradients, maximize, weight_decay, beta1, beta2, eps):
@@ -144,13 +136,13 @@ class Adamax(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)
         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),
+        self.hyper_map(ops.partial(_adamax_opt, beta1, beta2, eps, clr),
                        params, grads, exp_avg, exp_inf)
         return True
 

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