mindspore 2.4.0__cp310-cp310-macosx_10_15_x86_64.whl

mindspore/nn/optim/rprop.py

@@ -0,0 +1,251 @@
+# 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.
+# ============================================================================
+"""rprop"""
+from __future__ import absolute_import
+
+from mindspore import ops
+from mindspore.ops import operations as P
+import mindspore.common.dtype as mstype
+from mindspore.common.api import jit
+from mindspore import _checkparam as validator
+from mindspore.nn.optim.optimizer import Optimizer
+from mindspore.nn.optim.optimizer import opt_init_args_register
+
+
+class Rprop(Optimizer):
+    r"""
+    Implements Resilient backpropagation.
+
+    Further information about this implementation can be found at  `A Direct Adaptive Method for Faster Backpropagation
+    Learning: The RPROP Algorithm <https://ieeexplore.ieee.org/document/298623>`_.
+
+    The updating formulas are as follows:
+
+    .. math::
+        \begin{gather*}
+            &\hspace{-10mm}  \textbf{if} \:   g_{t-1} g_t  > 0                                     \\
+            &\hspace{25mm}  \Delta_t \leftarrow \mathrm{min}(\Delta_{t-1} \eta_{+}, \Delta_{max})  \\
+            &\hspace{0mm}  \textbf{else if}  \:  g_{t-1} g_t < 0                                   \\
+            &\hspace{25mm}  \Delta_t \leftarrow \mathrm{max}(\Delta_{t-1} \eta_{-}, \Delta_{min})  \\
+            &\hspace{-25mm}  \textbf{else}  \:                                                     \\
+            &\hspace{-5mm}  \Delta_t \leftarrow \Delta_{t-1}                                       \\
+            &\hspace{15mm} w_{t} \leftarrow w_{t-1}- \Delta_{t} \mathrm{sign}(g_t)                 \\
+        \end{gather*}
+
+    :math:`\Delta_{min/max}` represents the min/max step size, :math:`\eta_{+/-}` represents the factors of
+    etaminus and etaplus, :math:`g` represents `gradients`, :math:`w` represents `parameters`.
+
+    Note:
+        If parameters are not grouped, the `weight_decay` in optimizer will be applied on the parameters without 'beta'
+        or 'gamma' in their names.
+
+        Users can group parameters to change the strategy of decaying weight.
+
+        When parameters are grouped, each group can set `weight_decay`. If not, the `weight_decay` in optimizer will be
+        applied.
+
+    Args:
+        params (Union[list[Parameter], list[dict]]): Must be list of `Parameter` or list of `dict`. When the
+            `parameters` is a list of `dict`, the "params", "lr", "weight_decay", "grad_centralization" and
+            "order_params" are the keys can be parsed.
+
+            - params: Required. Parameters in current group. The value must be a list of `Parameter`.
+
+            - lr: Optional. If "lr" in the keys, the value of corresponding learning rate will be used.
+              If not, the `learning_rate` in optimizer will be used. Fixed and dynamic learning rate are supported.
+
+            - weight_decay: Optional. If "weight_decay" in the keys, the value of corresponding weight decay
+              will be used. If not, the `weight_decay` in the optimizer will be used. It should be noted that weight
+              decay can be a constant value or a Cell. It is a Cell only when dynamic weight decay is applied. Dynamic
+              weight decay is similar to dynamic learning rate, users need to customize a weight decay schedule only
+              with global step as input, and during training, the optimizer calls the instance of WeightDecaySchedule
+              to get the weight decay value of current step.
+
+            - grad_centralization: Optional. Must be Boolean. If "grad_centralization" is in the keys, the set value
+              will be used. If not, the `grad_centralization` is False by default. This configuration only works on the
+              convolution layer.
+
+            - order_params: Optional. When parameters is grouped, this usually is used to maintain the order of
+              parameters that appeared in the network to improve performance. The value should be parameters whose
+              order will be followed in optimizer.
+              If `order_params` in the keys, other keys will be ignored and the element of 'order_params' must be in
+              one group of `params`.
+
+        learning_rate (Union[float, int, Tensor, Iterable, LearningRateSchedule]): Learning_rate. Default: ``0.1`` .
+
+            - float: The fixed learning rate value. Must be equal to or greater than 0.
+
+            - int: The fixed learning rate value. Must be equal to or greater than 0. It will be converted to float.
+
+            - Tensor: Its value should be a scalar or a 1-D vector. For scalar, fixed learning rate will be applied.
+              For vector, learning rate is dynamic, then the i-th step will take the i-th value as the learning rate.
+
+            - Iterable: Learning rate is dynamic. The i-th step will take the i-th value as the learning rate.
+
+            - LearningRateSchedule: Learning rate is dynamic. During training, the optimizer calls the instance of
+              `LearningRateSchedule
+              <https://www.mindspore.cn/docs/en/master/api_python/mindspore.nn.html#learningrateschedule-class>`_
+              with step as the input to get the learning rate of current step.
+
+        etas (tuple[float, float]): The factor of multiplicative increasing or
+            descreasing(etaminus, etaplus). Default: ``(0.5, 1.2)`` .
+        step_sizes(tuple[float, float]): The allowed minimal and maximal step size(min_step_sizes, max_step_size).
+            Default: ``(1e-6, 50.)`` .
+        weight_decay (Union[float, int, Cell]): Weight decay (L2 penalty). Default: ``0.0`` .
+
+            - float: The fixed weight decay value. Must be equal to or greater than 0.
+
+            - int: The fixed weight decay value. Must be equal to or greater than 0. It will be converted to float.
+
+            - Cell: Weight decay is dynamic. During training, the optimizer calls the instance of
+              the Cell with step as the input to get the weight decay value of current step.
+
+    Inputs:
+        - **gradients** (tuple[Tensor]) - The gradients of `params`, the shape is the same as `params`.
+
+    Outputs:
+        Tensor[bool], the value is True.
+
+    Raises:
+        TypeError: If `learning_rate` is not one of int, float, Tensor, Iterable, LearningRateSchedule.
+        TypeError: If element of `parameters` is neither Parameter nor dict.
+        TypeError: If `step_sizes` or `etas` is not a tuple.
+        ValueError: If maximal step size is less than minimal step size.
+        ValueError: If the length of `step_sizes` or `etas` is not equal to 2.
+        TypeError: If  the element in `etas` or `step_sizes` is not a float.
+        ValueError: If `etaminus` is not in the range of (0, 1) or `etaplus` is not greater than 1.
+        TypeError: If `weight_decay` is neither float nor int.
+        ValueError: If `weight_decay` is less than 0.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import nn
+        >>>
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> #1) All parameters use the same learning rate and weight decay
+        >>> optim = nn.Rprop(params=net.trainable_params())
+        >>>
+        >>> #2) Use parameter groups and set different values
+        >>> conv_params = list(filter(lambda x: 'conv' in x.name, net.trainable_params()))
+        >>> no_conv_params = list(filter(lambda x: 'conv' not in x.name, net.trainable_params()))
+        >>> group_params = [{'params': conv_params,'grad_centralization':True},
+        ...                 {'params': no_conv_params, 'lr': 0.01},
+        ...                 {'order_params': net.trainable_params()}]
+        >>> optim = nn.Rprop(group_params, learning_rate=0.1, weight_decay=0.0)
+        >>> # The conv_params's parameters will use default learning rate of 0.1 default weight decay of 0.0 and grad
+        >>> # centralization of True.
+        >>> # The no_conv_params's parameters will use learning rate of 0.01 and default weight decay of 0.0 and grad
+        >>> # centralization of False.
+        >>> # The final parameters order in which the optimizer will be followed is the value of 'order_params'.
+        >>>
+        >>> loss = nn.SoftmaxCrossEntropyWithLogits()
+        >>> model = ms.train.Model(net, loss_fn=loss, optimizer=optim)
+    """
+
+    @opt_init_args_register
+    def __init__(self, params, learning_rate=0.1, etas=(0.5, 1.2), step_sizes=(1e-6, 50.), weight_decay=0.):
+
+        super(Rprop, self).__init__(learning_rate, params, weight_decay)
+        if not isinstance(etas, tuple):
+            raise TypeError("For Rprop, etas must be a tuple, but got {}.".format(type(etas)))
+        if len(etas) != 2:
+            raise ValueError("For Rprop, etas must be a tuple with the size of 2, but got {}.".format(len(etas)))
+
+        if not isinstance(step_sizes, tuple):
+            raise TypeError("For Rprop, step_sizes must be a tuple, but got {}.".format(type(etas)))
+        if len(step_sizes) != 2:
+            raise ValueError("For Rprop, step_sizes must be a tuple with the size of 2, "
+                             "but got {}.".format(len(step_sizes)))
+
+        if step_sizes[0] > step_sizes[1]:
+            raise ValueError("For Rprop, maximal step size should not be less than minimal step size, "
+                             "but got {} > {}.".format(step_sizes[0], step_sizes[1]))
+
+        validator.check_float_range(etas[0], 0.0, 1.0, validator.INC_NEITHER, "etaminus", self.cls_name)
+        validator.check_value_type("etaplus", etas[1], [float], self.cls_name)
+        if etas[1] <= 1.0:
+            raise ValueError("For Rprop, etaplus must be greater than 1.0, but got etaplus {}.".format(etas[1]))
+
+        validator.check_value_type("min_step_sizes", step_sizes[0], [float], self.cls_name)
+        validator.check_value_type("max_step_sizes", step_sizes[1], [float], self.cls_name)
+
+        self.etaminus, self.etaplus = etas
+        self.step_size_min, self.step_size_max = step_sizes
+        self.prev = self._parameters.clone(prefix="prev", init='zeros')
+        self.step_size = self._parameters.clone(prefix="step_size", init='zeros')
+
+        self.sign = P.Sign()
+        self.fill = P.FillV2()
+        self.assign = P.Assign()
+        self.assignadd = P.AssignAdd()
+        self.cast = P.Cast()
+        self.select = P.Select()
+        self.ones_like = P.OnesLike()
+
+    @jit
+    def construct(self, gradients):
+        gradients = self.flatten_gradients(gradients)
+        gradients = self.decay_weight(gradients)
+        gradients = self.gradients_centralization(gradients)
+        gradients = self.scale_grad(gradients)
+        lrs = self.get_lr()
+        self.assignadd(self.global_step, self.global_step_increase_tensor)
+        success = True
+
+        for index, (grad, param, prev, step_size) in enumerate(zip(gradients, self._parameters,
+                                                                   self.prev, self.step_size)):
+            lr = lrs[index] if self.is_group_lr else lrs
+
+            if self.global_step == 1:
+                step_size_fp32 = self.ones_like(step_size) * lr
+            else:
+                step_size_fp32 = self.cast(step_size, mstype.float32)
+
+            gradient_fp32 = self.cast(grad, mstype.float32)
+            param_fp32 = self.cast(param, mstype.float32)
+
+            sign = self.sign(gradient_fp32 * prev)
+            sign = self.select(
+                sign > 0,
+                self.fill(sign.shape, self.cast(self.etaplus, mstype.float32)),
+                sign)
+            sign = self.select(
+                sign < 0,
+                self.fill(sign.shape, self.cast(self.etaminus,
+                                                mstype.float32)), sign)
+            sign = self.select(
+                sign == 0, self.fill(sign.shape,
+                                     self.cast(1., mstype.float32)), sign)
+
+            step_size_fp32 = ops.clip_by_value(step_size_fp32 * sign,
+                                               self.step_size_min,
+                                               self.step_size_max)
+
+            gradient_update = self.select(
+                sign == self.etaminus,
+                self.fill(sign.shape, self.cast(0., mstype.float32)),
+                gradient_fp32)
+            next_param = param_fp32 - self.sign(gradient_update) * step_size_fp32
+
+            self.assign(param, self.cast(next_param, param.dtype))
+            self.assign(prev, self.cast(gradient_update, prev.dtype))
+            self.assign(step_size, self.cast(step_size_fp32, step_size.dtype))
+
+        return success

mindspore/nn/optim/sgd.py

@@ -0,0 +1,237 @@
+# Copyright 2020-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.
+# ============================================================================
+"""sgd"""
+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
+from mindspore.common.api import jit
+import mindspore.common.dtype as mstype
+from mindspore import _checkparam as validator
+from mindspore.nn.optim.optimizer import Optimizer
+from mindspore.nn.optim.optimizer import opt_init_args_register
+
+_sgd_opt = C.MultitypeFuncGraph("sgd_opt")
+
+
+@_sgd_opt.register("Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Function")
+def _tensor_run_opt_ext(momentum, learning_rate, gradient, weight, accum, stat, opt):
+    """Apply sgd optimizer to the weight parameter using Tensor."""
+    success = True
+    success = F.depend(success, opt(weight, gradient, learning_rate, accum, momentum, stat))
+    return success
+
+
+class SGD(Optimizer):
+    r"""
+    Implements stochastic gradient descent. Momentum is optional.
+
+    Introduction to SGD can be found at `SGD <https://en.wikipedia.org/wiki/Stochastic_gradient_dencent>`_ .
+    Nesterov momentum is based on the formula from paper `On the importance of initialization and
+    momentum in deep learning <http://proceedings.mlr.press/v28/sutskever13.html>`_.
+
+    .. math::
+        v_{t+1} = u \ast v_{t} + gradient \ast (1-dampening)
+
+    If nesterov is True:
+
+    .. math::
+        p_{t+1} = p_{t} - lr \ast (gradient + u \ast v_{t+1})
+
+    If nesterov is False:
+
+    .. math::
+        p_{t+1} = p_{t} - lr \ast v_{t+1}
+
+    To be noticed, for the first step, :math:`v_{t+1} = gradient`.
+
+    Here : where p, v and u denote the parameters, accum, and momentum respectively.
+
+    Note:
+        If parameters are not grouped, the `weight_decay` in optimizer will be applied on the network parameters without
+        'beta' or 'gamma' in their names. Users can group parameters to change the strategy of decaying weight. When
+        parameters are grouped, each group can set `weight_decay`. If not, the `weight_decay` in optimizer will be
+        applied.
+
+    Args:
+        params (Union[list[Parameter], list[dict]]): Must be list of `Parameter` or list of `dict`. When the
+            `params` is a list of `dict`, the string "params", "lr", "grad_centralization" and
+            "order_params" are the keys can be parsed.
+
+            - params: Required. Parameters in current group. The value must be a list of `Parameter`.
+
+            - lr: Optional. If "lr" in the keys, the value of corresponding learning rate will be used.
+              If not, the `learning_rate` in optimizer will be used. Fixed and dynamic learning rate are supported.
+
+            - weight_decay: Optional. If "weight_decay" in the keys, the value of corresponding weight decay
+              will be used. If not, the `weight_decay` in the optimizer will be used. It should be noted that weight
+              decay must be float, dynamic weight decay is currently not supported.
+
+            - grad_centralization: Optional. Must be Boolean. If "grad_centralization" is in the keys, the set value
+              will be used. If not, the `grad_centralization` is False by default. This configuration only works on the
+              convolution layer.
+
+            - order_params: Optional. When parameters is grouped, this usually is used to maintain the order of
+              parameters that appeared in the network to improve performance. The value should be parameters whose
+              order will be followed in optimizer.
+              If `order_params` in the keys, other keys will be ignored and the element of 'order_params' must be in
+              one group of `params`.
+
+        learning_rate (Union[float, int, Tensor, Iterable, LearningRateSchedule]): Default: ``0.1`` .
+
+            - float: The fixed learning rate value. Must be equal to or greater than 0.
+
+            - int: The fixed learning rate value. Must be equal to or greater than 0. It will be converted to float.
+
+            - Tensor: Its value should be a scalar or a 1-D vector. For scalar, fixed learning rate will be applied.
+              For vector, learning rate is dynamic, then the i-th step will take the i-th value as the learning rate.
+
+            - Iterable: Learning rate is dynamic. The i-th step will take the i-th value as the learning rate.
+
+            - LearningRateSchedule: Learning rate is dynamic. During training, the optimizer calls the instance of
+              `LearningRateSchedule
+              <https://www.mindspore.cn/docs/en/master/api_python/mindspore.nn.html#learningrateschedule-class>`_
+              with step as the input to get the learning rate of current step.
+
+        momentum (float): A floating point value the momentum. must be at least 0.0. Default: ``0.0`` .
+        dampening (float): A floating point value of dampening for momentum. must be at least 0.0. Default: ``0.0`` .
+        weight_decay (float): Weight decay (L2 penalty). It must be equal to or greater than 0. Default: ``0.0`` .
+        nesterov (bool): Enables the Nesterov momentum. If use nesterov, momentum must be positive,
+                         and dampening must be equal to 0.0. Default: ``False`` .
+        loss_scale (float): A floating point value for the loss scale, which must be larger than 0.0. In general, use
+            the default value. Only when `FixedLossScaleManager` is used for training and the `drop_overflow_update` in
+            `FixedLossScaleManager` is set to ``False`` , then this value needs to be the same as the `loss_scale` in
+            `FixedLossScaleManager`. Refer to class :class:`mindspore.amp.FixedLossScaleManager` for more details.
+            Default: ``1.0`` .
+
+    Inputs:
+        - **gradients** (tuple[Tensor]) - The gradients of `params`, the shape is the same as `params`.
+
+    Outputs:
+        Tensor[bool], the value is ``True`` .
+
+    Raises:
+        ValueError: If the momentum, dampening or weight_decay value is less than 0.0.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import nn
+        >>>
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> #1) All parameters use the same learning rate and weight decay
+        >>> optim = nn.SGD(params=net.trainable_params())
+        >>>
+        >>> #2) Use parameter groups and set different values
+        >>> conv_params = list(filter(lambda x: 'conv' in x.name, net.trainable_params()))
+        >>> no_conv_params = list(filter(lambda x: 'conv' not in x.name, net.trainable_params()))
+        >>> group_params = [{'params': conv_params,'grad_centralization':True},
+        ...                 {'params': no_conv_params, 'lr': 0.01},
+        ...                 {'order_params': net.trainable_params()}]
+        >>> optim = nn.SGD(group_params, learning_rate=0.1, weight_decay=0.0)
+        >>> # The conv_params's parameters will use default learning rate of 0.1 and default weight decay of 0.0
+        >>> # and grad centralization of True.
+        >>> # The no_conv_params's parameters will use learning rate of 0.01 and default weight decay of 0.0 and grad
+        >>> # centralization of False.
+        >>> # The final parameters order in which the optimizer will be followed is the value of 'order_params'.
+        >>>
+        >>> loss = nn.SoftmaxCrossEntropyWithLogits()
+        >>> model = ms.train.Model(net, loss_fn=loss, optimizer=optim)
+    """
+
+    @opt_init_args_register
+    def __init__(self, params, learning_rate=0.1, momentum=0.0, dampening=0.0, weight_decay=0.0, nesterov=False,
+                 loss_scale=1.0):
+
+        super(SGD, self).__init__(learning_rate, params, weight_decay, loss_scale)
+
+        if isinstance(momentum, int):
+            momentum = float(momentum)
+        if not isinstance(momentum, float):
+            raise TypeError(f"For 'SGD', the argument 'momentum' must be float type, "
+                            f"but got {type(momentum)}.")
+
+        if isinstance(momentum, float) and momentum < 0.0:
+            raise ValueError(f"For 'SGD', the argument 'momentum' must be at least 0.0, "
+                             f"but got {momentum}.")
+
+        if isinstance(dampening, int):
+            dampening = float(dampening)
+        if not isinstance(dampening, float):
+            raise TypeError(f"For 'SGD', the argument 'dampening' must be float type, "
+                            f"but got {type(dampening)}.")
+
+        if dampening < 0.0:
+            raise ValueError(f"For 'SGD', the argument 'dampening' must be at least 0.0, "
+                             f"but got 'dampening' {dampening}")
+        self.dampening = dampening
+
+        validator.check_value_type("nesterov", nesterov, [bool], self.cls_name)
+
+        if nesterov and (momentum <= 0.0 or dampening != 0.0):
+            raise ValueError(f"For 'SGD', if 'nesterov' is true, 'momentum' must be > 0.0 and 'dampening' must "
+                             f"equal to 0.0, but got 'momentum' {momentum}, 'dampening' {dampening}.")
+        self.nesterov = nesterov
+
+        if self.dynamic_weight_decay:
+            raise TypeError("For 'SGD', dynamic weight decay is currently not supported, the argument 'weight_decay' "
+                            "or 'weight_decay' set in grouped 'params' must be float or int type.")
+
+        if hasattr(self, "group_weight_decay") and self.group_weight_decay:
+            self.opt = tuple(P.SGD(dampening, 0.0, nesterov) for wd in self.group_weight_decay)
+        else:
+            self.opt = tuple([P.SGD(dampening, 0.0, nesterov)] * len(self._parameters))
+
+        self.momentum = Parameter(Tensor(momentum, mstype.float32), name="momentum")
+
+        if not momentum > 0.0:
+            enable_cache_param_list = []
+            for param in self._parameters:
+                if param.cache_enable:
+                    enable_cache_param_list.append(param)
+                    param.cache_enable = False
+
+        self.accum = self._parameters.clone(prefix="accum", init='zeros')
+        self.stat = self._parameters.clone(prefix="stat", init='ones')
+
+
+        if not momentum > 0.0:
+            for param in enable_cache_param_list:
+                param.cache_enable = True
+
+
+    @jit
+    def construct(self, gradients):
+        params = self._parameters
+        accum = self.accum
+        stat = self.stat
+        gradients = self.decay_weight(gradients)
+        gradients = self.flatten_gradients(gradients)
+        gradients = self.gradients_centralization(gradients)
+        gradients = self.scale_grad(gradients)
+        lr = self.get_lr()
+        self.assignadd(self.global_step, self.global_step_increase_tensor)
+        if self.is_group_lr:
+            success = self.hyper_map_reverse(F.partial(_sgd_opt, self.momentum),
+                                             lr, gradients, params, accum, stat, self.opt)
+        else:
+            success = self.hyper_map_reverse(F.partial(_sgd_opt, self.momentum, lr),
+                                             gradients, params, accum, stat, self.opt)
+        return success

mindspore/nn/optim/tft_wrapper.py

@@ -0,0 +1,127 @@
+# Copyright 2024 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.
+# ============================================================================
+"""OptTFTWrapper"""
+from __future__ import absolute_import
+
+import os
+from mindspore.common.tensor import Tensor
+from mindspore.nn.optim.optimizer import Optimizer
+from mindspore.ops.operations.manually_defined._inner import TensorReport
+from mindspore import ops, context
+
+class OptTFTWrapper(Optimizer):
+    r"""
+    Implements TFT optimizer wrapper, this wrapper is used to report status to MindIO TFT before optimizer updating.
+
+    Note:
+        This optimizer is depend on MindIO TFT feature. Currently only support ascend graph mode and
+        sink_size must be less than 1.
+
+    Args:
+        opt (Optimizer): Must be sub-class of Optimizer.
+
+    Inputs:
+        - **gradients** (tuple[Tensor]) - The gradients of opt's `params`, the shape is the same as opt's `params`.
+
+    Outputs:
+        Tensor, result of executing optimizer 'opt'.
+
+    Raises:
+        TypeError: If the parameter opt is not an subclass of Optimizer.
+        ValueError: If the platform is not Ascend graph mode, or customer doesn't switch on TFT feature.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import nn
+        >>>
+        >>> # Define the network structure of LeNet5. Refer to
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
+        >>> net = LeNet5()
+        >>> #1) All parameters use the same learning rate and weight decay
+        >>> optim = nn.SGD(params=net.trainable_params())
+        >>> optim_wrapper = nn.OptTFTWrapper(optim)
+        >>>
+        >>> loss = nn.SoftmaxCrossEntropyWithLogits()
+        >>> model = ms.train.Model(net, loss_fn=loss, optimizer=optim)
+    """
+
+    def __init__(self, opt, **kwargs):
+        super(OptTFTWrapper, self).__init__(opt.learning_rate, opt._parameters) # pylint: disable=W0212
+        if not isinstance(opt, Optimizer):
+            raise TypeError(f"For 'OptTFTWrapper', the argument 'opt' must be Optimizer type, " f"but got {type(opt)}.")
+        tft_env = os.getenv("MS_ENABLE_TFT", "")
+        if ("TTP:1" not in tft_env) and ("UCE:1" not in tft_env):
+            raise ValueError("MindIO TFT regitster need custom switch on[MS_ENABLE_TFT='{TTP:1,UCE:1}']!")
+        mode = context.get_context("mode")
+        device_target = context.get_context("device_target")
+        if device_target != "Ascend" or mode != context.GRAPH_MODE:
+            raise ValueError("MindIO adataper only support on Ascend device with GRAPH Mode!")
+        self.opt = opt
+        self.report = TensorReport()
+        self.depend = ops.Depend()
+        self.g_one = Tensor([0.1])
+        # enable consistent check by default, only disable when enable_consistent_check is False
+        self.use_allreduce = kwargs.get("enable_consistent_check", True)
+
+        if self.use_allreduce:
+            self.allreduce_sum = ops.AllReduce()
+            self.allreduce_sum.add_prim_attr("tft_report_before", True)
+
+        self.param_rank = opt.param_rank
+        self.optim_filter = opt.optim_filter
+        self.loss_scale = opt.loss_scale
+        self.dynamic_weight_decay = opt.dynamic_weight_decay
+        self.grad_centralization = opt.grad_centralization
+
+        self.dynamic_lr = opt.dynamic_lr
+        self.global_step = opt.global_step
+        self.is_group = opt.is_group
+        self.is_group_lr = opt.is_group_lr
+        self.is_group_params_ordered = opt.is_group_params_ordered
+        self.use_parallel = opt.use_parallel
+        if self.is_group:
+            self.group_params = opt.group_params
+            self.group_lr = opt.group_lr
+            self.group_weight_decay = opt.group_weight_decay
+            self.group_grad_centralization = opt.group_grad_centralization
+            self.grad_centralization_flags = opt.grad_centralization_flags
+
+        self.skip_auto_parallel_compile = opt.skip_auto_parallel_compile
+
+        self.learning_rate = opt.learning_rate
+        self.parameters = opt.parameters
+        self.decay_flags = opt.decay_flags
+        self.dynamic_decay_flags = opt.dynamic_decay_flags
+        self.weight_decay = opt.weight_decay
+        self.exec_weight_decay = opt.exec_weight_decay
+        self.ps_parameters = opt.ps_parameters
+        self.cache_enable = opt.cache_enable
+        self.reciprocal_scale = opt.reciprocal_scale
+        self.need_scale = opt.need_scale
+        self.global_step_increase_tensor = opt.global_step_increase_tensor
+        self.param_length = opt.param_length
+        self.enable_tuple_broaden = opt.enable_tuple_broaden
+
+    def construct(self, gradients):
+        g_one = self.depend(self.g_one, gradients)
+        if self.use_allreduce is True:
+            g_one_res = self.allreduce_sum(g_one)
+        else:
+            g_one_res = g_one
+        self.report("tft_report", g_one_res)
+        return self.opt(gradients)