mindspore 2.4.10__cp39-cp39-win_amd64.whl → 2.6.0__cp39-cp39-win_amd64.whl

mindspore/mint/nn/layer/pooling.py

@@ -104,7 +104,242 @@ class AdaptiveAvgPool2d(_AdaptiveAvgPoolNd):
         return mint.nn.functional.adaptive_avg_pool2d(input, self.output_size)
 
 
+class AdaptiveAvgPool3d(Cell):
+    r"""
+    This operator applies a 3D adaptive average pooling to an input signal composed of multiple input planes.
+    That is, for any input size, the size of the specified output is :math:`(D, H, W)`.
+    The number of output features is equal to the number of input planes.
+
+    Suppose the last 3 dimension size of input is :math:`(inD, inH, inW)`, then the last 3 dimension size of output is
+    :math:`(outD, outH, outW)`.
+
+    .. math::
+        \begin{array}{ll} \\
+            \forall \quad od \in [0,outD-1], oh \in [0,outH-1], ow \in [0,outW-1]\\
+            output[od,oh,ow] = \\
+            \qquad mean(input[istartD:iendD+1,istartH:iendH+1,istartW:iendW+1])\\
+            where,\\
+            \qquad istartD= \left\lceil \frac{od * inD}{outD} \right\rceil \\
+            \qquad iendD=\left\lfloor \frac{(od+1)* inD}{outD} \right\rfloor \\
+            \qquad istartH=\left\lceil \frac{oh * inH}{outH} \right\rceil \\
+            \qquad iendH=\left\lfloor \frac{(oh+1) * inH}{outH} \right\rfloor \\
+            \qquad istartW=\left\lceil \frac{ow * inW}{outW} \right\rceil \\
+            \qquad iendW=\left\lfloor \frac{(ow+1) * inW}{outW} \right\rfloor
+        \end{array}
+
+    .. warning::
+        For Ascend, it is only supported on Atlas A2 Training Series Products.
+        This is an experimental optimizer API that is subject to change or deletion.
+
+    Args:
+        output_size (Union[int, tuple]): The target output size. `output_size` can be a tuple :math:`(D, H, W)`,
+            or an int D for :math:`(D, D, D)`. :math:`D`, :math:`H` and :math:`W` can be int or None
+            which means the output size is the same as that of the input.
+
+    Inputs:
+        - **input** (Tensor) - The input of AdaptiveAvgPool3d, which is a 5D or 4D Tensor.
+
+    Outputs:
+        Tensor, with the same type as the `input`.
+
+    Raises:
+        TypeError: If `input` is not a Tensor.
+        ValueError: If the dimension of `input` is not 4D or 5D.
+        ValueError: If `output_size` value is not positive.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> from mindspore import mint
+        >>> import numpy as np
+        >>> # case 1: output_size=(3, 3, 4)
+        >>> output_size=(3, 3, 4)
+        >>> input_x_val = np.random.randn(4, 3, 5, 6, 7)
+        >>> input_x = ms.Tensor(input_x_val, ms.float32)
+        >>> net = mint.nn.AdaptiveAvgPool3d(output_size)
+        >>> output = net(input_x)
+        >>> print(output.shape)
+        (4, 3, 3, 3, 4)
+        >>> # case 2: output_size=4
+        >>> output_size=5
+        >>> input_x_val = np.random.randn(2, 3, 8, 6, 12)
+        >>> input_x = ms.Tensor(input_x_val, ms.float32)
+        >>> net = mint.nn.AdaptiveAvgPool3d(output_size)
+        >>> output = net(input_x)
+        >>> print(output.shape)
+        (2, 3, 5, 5, 5)
+        >>> # case 3: output_size=(None, 4, 5)
+        >>> output_size=(None, 4, 5)
+        >>> input_x_val = np.random.randn(4, 1, 9, 10, 8)
+        >>> input_x = ms.Tensor(input_x_val, ms.float32)
+        >>> net = mint.nn.AdaptiveAvgPool3d(output_size)
+        >>> output = net(input_x)
+        >>> print(output.shape)
+        (4, 1, 9, 4, 5)
+    """
+
+    def __init__(self, output_size):
+        """Initialize AdaptiveAvgPool3d."""
+        super(AdaptiveAvgPool3d, self).__init__()
+        self.output_size = output_size
+
+    def construct(self, input):
+        return mint.nn.functional.adaptive_avg_pool3d(input, self.output_size)
+
+
+class MaxUnpool2d(Cell):
+    r"""
+    Computes the inverse of `Maxpool2d`.
+
+    `MaxUnpool2d` keeps the maximal value and set all position of non-maximal values to zero.
+    Typically the input is of shape :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`,
+    and the output is of shape :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`.
+    The operation is as follows.
+
+    .. math::
+        \begin{array}{ll} \\
+        H_{out} = (H_{in} - 1) \times stride[0] - 2 \times padding[0] + kernel\_size[0] \\
+        W_{out} = (W_{in} - 1) \times stride[1] - 2 \times padding[1] + kernel\_size[1] \\
+        \end{array}
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        kernel_size (Union[int, tuple[int]]): The size of kernel used to take the maximum value,
+            an int number that represents height and width of the kernel,
+            or a tuple of two int numbers that represent height and width respectively.
+        stride (Union[int, tuple[int]], optional): The distance of kernel moving,
+            an int number that represents the height and width of movement are both stride,
+            or a tuple of two int numbers that represent height and width of movement respectively.
+            Default: ``None`` , which indicates the moving step is `kernel_size` .
+        padding (Union[int, tuple[int]], optional): The pad value to be filled. Default: ``0`` .
+            If `padding` is an integer, the paddings of height and width are the same, equal to padding.
+            If `padding` is a tuple of two integers, the padding of height and width equal to padding[0]
+            and padding[1] correspondingly.
+
+    Inputs:
+        - **input** (Tensor) - The input Tensor to invert.
+          Tensor of shape :math:`(N, C, H_{in}, W_{in})` or :math:`(C, H_{in}, W_{in})`.
+        - **indices** (Tensor) - Max values' index represented by the indices.
+          Tensor of shape must be same with input 'input'.
+          Values of indices must belong to :math:`[0, H_{in} \times W_{in} - 1]`.
+          Data type must be in int32 or int64.
+        - **output_size** (tuple[int], optional) - The target output size. Default: ``None`` .
+          If output_size == (), then the shape of output computed by `kernel_size`, `stride` and `padding`.
+          If output_size != (), then output_size must be :math:`(N, C, H, W)` , :math:`(C, H, W)` or :math:`(H, W)`
+          and output_size must belong to
+          :math:`[(N, C, H_{out} - stride[0], W_{out} - stride[1]), (N, C, H_{out} + stride[0], W_{out} + stride[1])]`.
+
+    Outputs:
+        Tensor, with shape :math:`(N, C, H_{out}, W_{out})` or :math:`(C, H_{out}, W_{out})`,
+        with the same data type with `input`.
+
+    Raises:
+        TypeError: If data type of `input` or `indices` is not supported.
+        TypeError: If `kernel_size`, `stride` or `padding` is neither an int nor a tuple.
+        ValueError: If numbers in `stride`, `padding` or `kernel_size` is not positive.
+        ValueError: If the shapes of `input` and `indices` are not equal.
+        ValueError: If `input` whose length is not 3 or 4.
+        ValueError: If `output_size` whose type is not tuple.
+        ValueError: If `output_size` is not close to output size computed by attr `kernel_size`, `stride`, `padding`.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import numpy as np
+        >>> from mindspore import Tensor, mint
+        >>> input = Tensor(np.array([[[[0, 1], [8, 9]]]]).astype(np.float32))
+        >>> indices = Tensor(np.array([[[[0, 1], [2, 3]]]]).astype(np.int64))
+        >>> net =  mint.nn.MaxUnpool2d(1, stride=1, padding=0)
+        >>> output = net(input, indices)
+        >>> print(output.asnumpy())
+        [[[[0. 1.]
+           [8. 9.]]]]
+    """
+
+    def __init__(self, kernel_size, stride=None, padding=0) -> None:
+        super(MaxUnpool2d, self).__init__()
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+
+    def construct(self, input, indices, output_size=None):
+        return mint.nn.functional.max_unpool2d(input, indices,
+                                               self.kernel_size, self.stride,
+                                               self.padding, output_size)
+
+
+class _AdaptiveMaxPoolNd(Cell):
+    """Common base of AdaptiveMaxPool1d"""
+
+    def __init__(self, output_size, return_indices=False) -> None:
+        super(_AdaptiveMaxPoolNd, self).__init__()
+        self.output_size = output_size
+        self.return_indices = return_indices
+
+    def extend_repr(self):
+        return 'output_size={}, return_indices={}'.format(self.output_size, self.return_indices)
+
+
+class AdaptiveMaxPool1d(_AdaptiveMaxPoolNd):
+    r"""
+    Applies a 1D adaptive max pooling over an input signal composed of several input planes.
+
+    The output is of size :math:`L_{out}` , for any input size.
+    The number of output features is equal to the number of input planes.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        output_size (Union[int, tuple]): the target output size :math:`L_{out}` .
+        return_indices (bool, optional): Whether to return the index of the maximum value. Default: ``False`` .
+
+    Inputs:
+        - **input** (Tensor) - The input with shape :math:`(N, C, L_{in})` or :math:`(C, L_{in})` .
+
+    Outputs:
+        Union(Tensor, tuple(Tensor, Tensor)).
+
+        - If `return_indices` is False, output is a Tensor, with shape :math:`(N, C, L_{out})`. It has the same data
+          type as `x`.
+        - If `return_indices` is True, output is a Tuple of 2 Tensors, representing the result and where the max
+          values are generated.
+
+    Raises:
+        TypeError: If `input` is not a tensor.
+        TypeError: If dtype of `input` is not float16, float32 or float64.
+        TypeError: If `output_size` is not int or tuple.
+        TypeError: If `return_indices` is not a bool.
+        ValueError: If `output_size` is a tuple and the length of `output_size` is not 1.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> import numpy as np
+        >>> input = Tensor(np.array([[[2, 1, 2], [2, 3, 5]]]), mindspore.float16)
+        >>> net = mint.nn.AdaptiveMaxPool1d(3)
+        >>> output = net(input)
+        >>> print(output)
+        [[[2. 1. 2.]
+          [2. 3. 5.]]]
+    """
+
+    def construct(self, input):
+        return mint.nn.functional.adaptive_max_pool1d(input, self.output_size, self.return_indices)
+
+
 __all__ = [
+    'AdaptiveAvgPool3d',
     'AdaptiveAvgPool2d',
     'AdaptiveAvgPool1d',
+    'AdaptiveMaxPool1d',
+    'MaxUnpool2d',
 ]

mindspore/mint/optim/__init__.py

@@ -20,5 +20,7 @@ The optimizer is used to calculate and update the gradients.
 """
 from __future__ import absolute_import
 from mindspore.mint.optim.adamw import AdamW
+from mindspore.mint.optim.adam import Adam
+from mindspore.mint.optim.sgd import SGD
 
-__all__ = ['AdamW']
+__all__ = ['AdamW', 'Adam', 'SGD']

mindspore/mint/optim/adam.py

@@ -0,0 +1,223 @@
+# 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.
+# ============================================================================
+"""Adam"""
+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 import dtype as mstype
+from mindspore.experimental.optim.optimizer import Optimizer
+from mindspore import _checkparam as validator
+from mindspore import mint
+
+_optim_adamw_opt = C.MultitypeFuncGraph("optim_adamw_opt")
+hyper_map = C.HyperMap()
+assign_add = P.AssignAdd()
+
+
+@_optim_adamw_opt.register("Float", "Float", "Float", "Tensor", "Tensor", "Tensor", "Tensor",
+                           "Tensor", "Tensor", "Tensor")
+def _run_optim_adamw_amsgrad_opt(beta1, beta2, eps, neg_step_size, sqrt_bias_correction2, parameters, grads, exp_avg,
+                                 exp_avg_sq, max_exp_avg_sq):
+    """Apply adam optimizer to the weight parameter when amsgrad is True."""
+    success = True
+    exp_avg_tmp = mint.add(mint.mul(exp_avg, beta1), grads, alpha=1 - beta1)
+    exp_avg_sq_tmp = mint.mul(exp_avg_sq, beta2) + mint.mul(mint.mul(grads, grads), 1 - beta2)
+
+    max_exp_avg_sq = mint.maximum(max_exp_avg_sq, exp_avg_sq_tmp)
+    denom = F.cast(mint.div(mint.sqrt(max_exp_avg_sq), sqrt_bias_correction2), max_exp_avg_sq.dtype)
+    denom = mint.add(denom, eps)
+
+    delta_param = mint.mul(F.cast(neg_step_size, max_exp_avg_sq.dtype), mint.div(exp_avg_tmp, denom))
+    F.assign(exp_avg, exp_avg_tmp)
+    F.assign(exp_avg_sq, exp_avg_sq_tmp)
+    assign_add(parameters, delta_param)
+    return success
+
+
+@_optim_adamw_opt.register("Float", "Float", "Float", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor", "Tensor")
+def _run_optim_adamw_opt(beta1, beta2, eps, neg_step_size, sqrt_bias_correction2, parameters, grads, exp_avg,
+                         exp_avg_sq):
+    """Apply adam optimizer to the weight parameter when amsgrad is False."""
+    success = True
+    exp_avg_tmp = mint.add(mint.mul(exp_avg, beta1), grads, alpha=1 - beta1)
+    exp_avg_sq_tmp = mint.mul(exp_avg_sq, beta2) + mint.mul(mint.mul(grads, grads), 1 - beta2)
+
+    denom = F.cast(mint.div(mint.sqrt(exp_avg_sq_tmp), sqrt_bias_correction2), exp_avg_sq_tmp.dtype)
+    denom = mint.add(denom, eps)
+
+    delta_param = mint.mul(F.cast(neg_step_size, exp_avg_sq_tmp.dtype), mint.div(exp_avg_tmp, denom))
+    F.assign(exp_avg, exp_avg_tmp)
+    F.assign(exp_avg_sq, exp_avg_sq_tmp)
+    assign_add(parameters, delta_param)
+    return success
+
+
+def _check_param_value(betas, eps, weight_decay, lr, amsgrad, maximize, prim_name):
+    """Check the type of inputs."""
+    validator.check_value_type('betas', betas, [tuple], prim_name)
+    validator.check("betas size", len(betas), "", [2], validator.IN, prim_name)
+    validator.check_value_type("betas[0]", betas[0], [float], prim_name)
+    validator.check_value_type("betas[1]", betas[1], [float], prim_name)
+    validator.check_value_type("eps", eps, [float], prim_name)
+    validator.check_value_type("weight_decay", weight_decay, [float], prim_name)
+    validator.check_value_type("lr", lr, [float], prim_name)
+    validator.check_value_type("amsgrad", amsgrad, [bool], prim_name)
+    validator.check_value_type("maximize", maximize, [bool], prim_name)
+
+
+class Adam(Optimizer):
+    r"""
+    Implements Adaptive Moment Estimation (Adam) algorithm.
+
+    The updating formulas are as follows:
+
+    .. math::
+        \begin{aligned}
+            &\textbf{input}      : \gamma \text{ (lr)}, \beta_1, \beta_2
+                \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)}          \\
+            &\hspace{13mm}      \lambda \text{ (weight decay)},  \: \textit{amsgrad},
+                \:\textit{maximize}                                                              \\
+            &\textbf{initialize} :  m_0 \leftarrow 0 \text{ ( first moment)},
+                v_0\leftarrow 0 \text{ (second moment)},\: \widehat{v_0}^{max}\leftarrow 0\\[-1.ex]
+            &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
+            &\hspace{5mm}\textbf{if} \: \textit{maximize}:                                       \\
+            &\hspace{10mm}g_t           \leftarrow   -\nabla_{\theta} f_t (\theta_{t-1})         \\
+            &\hspace{5mm}\textbf{else}                                                           \\
+            &\hspace{10mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})          \\
+            &\hspace{5mm}\textbf{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}v_t           \leftarrow   \beta_2 v_{t-1} + (1-\beta_2) g^2_t          \\
+            &\hspace{5mm}\widehat{m_t} \leftarrow   m_t/\big(1-\beta_1^t \big)                   \\
+            &\hspace{5mm}\widehat{v_t} \leftarrow   v_t/\big(1-\beta_2^t \big)                   \\
+            &\hspace{5mm}\textbf{if} \: amsgrad                                                  \\
+            &\hspace{10mm}\widehat{v_t}^{max} \leftarrow \mathrm{max}(\widehat{v_t}^{max},
+                \widehat{v_t})                                                                   \\
+            &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/
+                \big(\sqrt{\widehat{v_t}^{max}} + \epsilon \big)                                 \\
+            &\hspace{5mm}\textbf{else}                                                           \\
+            &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/
+                \big(\sqrt{\widehat{v_t}} + \epsilon \big)                                       \\
+            &\bf{return} \:  \theta_t                                                     \\[-1.ex]
+       \end{aligned}
+
+    .. warning::
+            This is an experimental API that is subject to change or deletion.
+
+    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-3``.
+        betas (Tuple[float, float], optional): The exponential decay rate for the moment estimations.
+            Should be in range (0.0, 1.0). Default: ``(0.9, 0.999)``.
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. Should be greater than 0. Default: ``1e-8``.
+        weight_decay (float, optional): weight decay (L2 penalty). Default: ``0.``.
+        amsgrad (bool, optional): whether to use the AMSGrad algorithm. Default: ``False``.
+
+    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 `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` is not in the range of [0, 1).
+        ValueError: If the `weight_decay` is less than 0.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import nn
+        >>> from mindspore import mint
+        >>> # 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 = mint.optim.Adam(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-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0.0, amsgrad=False, *, maximize=False):
+        _check_param_value(betas, eps, weight_decay, lr, amsgrad, maximize, self.cls_name)
+        if lr < 0.0:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if eps < 0.0:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        if weight_decay < 0.0:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad,
+                        maximize=maximize)
+        self.max_v_group = True
+        super(Adam, self).__init__(params, defaults)
+
+        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.state_step = Parameter(Tensor([0], mstype.float32), "state_step")
+        self.increase_tensor = Tensor(1, mstype.float32)
+        self.assignadd = P.AssignAdd()
+        self.pow = P.Pow()
+
+
+    def construct(self, gradients):
+        self.assignadd(self.state_step, self.increase_tensor)
+        for group_id, group in enumerate(self.param_groups):
+            beta1, beta2 = group['betas']
+            maximize = group.get("maximize")
+            start_id = self.group_start_id[group_id]
+            end_id = self.group_start_id[group_id + 1]
+            lr = group.get("lr")
+            grads = tuple([grad if not maximize else mint.neg(grad) for grad in gradients[start_id: end_id]])
+
+            bias_correction1 = 1 - beta1 ** self.state_step
+            bias_correction2 = 1 - beta2 ** self.state_step
+            neg_step_size = -mint.div(lr, bias_correction1)
+            sqrt_bias_correction2 = mint.sqrt(bias_correction2)
+            grads = self._decay_weight(group.get("weight_decay"), self.parameters[start_id: end_id], grads)
+
+            if group.get("amsgrad"):
+                self.hyper_map(F.partial(_optim_adamw_opt, beta1, beta2, group.get("eps"), neg_step_size,
+                                         sqrt_bias_correction2),
+                               self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
+                               self.exp_avg_sq[start_id: end_id], group.get("max_exp_avg_sq"))
+            else:
+                self.hyper_map(F.partial(_optim_adamw_opt, beta1, beta2, group.get("eps"), neg_step_size,
+                                         sqrt_bias_correction2),
+                               self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
+                               self.exp_avg_sq[start_id: end_id])
+        return True

mindspore/mint/optim/adamw.py

@@ -15,7 +15,7 @@
 """adamw"""
 from __future__ import absolute_import
 
-from mindspore.ops import functional as F, composite as C, operations as P
+from mindspore.ops import functional as F, composite as C
 from mindspore.common.parameter import Parameter
 from mindspore.common.tensor import Tensor
 from mindspore.common import dtype as mstype
@@ -62,7 +62,9 @@ class AdamW(Optimizer):
     Implements Adam Weight Decay algorithm.
 
     .. math::
-        \begin{aligned}
+        \begin{array}{l}
+            &\newline
+            &\hline \\
             &\textbf{input}      : \gamma \text{(lr)}, \: \beta_1, \beta_2
                 \text{(betas)}, \: \theta_0 \text{(params)}, \: f(\theta) \text{(objective)},
                 \: \epsilon \text{ (epsilon)}                                                    \\
@@ -70,26 +72,32 @@ class AdamW(Optimizer):
                 \: \textit{maximize}                                                             \\
             &\textbf{initialize} : m_0 \leftarrow 0 \text{ (first moment)}, v_0 \leftarrow 0
                 \text{ ( second moment)}, \: \widehat{v_0}^{max}\leftarrow 0              \\[-1.ex]
+            &\newline
+            &\hline \\
             &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do}                         \\
-            &\hspace{5mm}\textbf{if} \: \textit{maximize}:                                       \\
-            &\hspace{10mm}g_t           \leftarrow   -\nabla_{\theta} f_t (\theta_{t-1})          \\
-            &\hspace{5mm}\textbf{else}                                                           \\
-            &\hspace{10mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})           \\
-            &\hspace{5mm} \theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1}         \\
-            &\hspace{5mm}m_t           \leftarrow   \beta_1 m_{t-1} + (1 - \beta_1) g_t          \\
-            &\hspace{5mm}v_t           \leftarrow   \beta_2 v_{t-1} + (1-\beta_2) g^2_t          \\
-            &\hspace{5mm}\widehat{m_t} \leftarrow   m_t/\big(1-\beta_1^t \big)                   \\
-            &\hspace{5mm}\widehat{v_t} \leftarrow   v_t/\big(1-\beta_2^t \big)                   \\
-            &\hspace{5mm}\textbf{if} \: amsgrad                                                  \\
-            &\hspace{10mm}\widehat{v_t}^{max} \leftarrow \mathrm{max}(\widehat{v_t}^{max},
+            &\hspace{6mm}\textbf{if} \: \textit{maximize}:                                       \\
+            &\hspace{11mm}g_t           \leftarrow   -\nabla_{\theta} f_t (\theta_{t-1})          \\
+            &\hspace{6mm}\textbf{else}                                                           \\
+            &\hspace{11mm}g_t           \leftarrow   \nabla_{\theta} f_t (\theta_{t-1})           \\
+            &\hspace{6mm} \theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1}         \\
+            &\hspace{6mm}m_t           \leftarrow   \beta_1 m_{t-1} + (1 - \beta_1) g_t          \\
+            &\hspace{6mm}v_t           \leftarrow   \beta_2 v_{t-1} + (1-\beta_2) g^2_t          \\
+            &\hspace{6mm}\widehat{m_t} \leftarrow   m_t/\big(1-\beta_1^t \big)                   \\
+            &\hspace{6mm}\widehat{v_t} \leftarrow   v_t/\big(1-\beta_2^t \big)                   \\
+            &\hspace{6mm}\textbf{if} \: amsgrad                                                  \\
+            &\hspace{11mm}\widehat{v_t}^{max} \leftarrow \mathrm{max}(\widehat{v_t}^{max},
                 \widehat{v_t})                                                                   \\
-            &\hspace{10mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/
+            &\hspace{11mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/
                 \big(\sqrt{\widehat{v_t}^{max}} + \epsilon \big)                                 \\
-            &\hspace{5mm}\textbf{else}                                                           \\
-            &\hspace{10mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/
+            &\hspace{6mm}\textbf{else}                                                           \\
+            &\hspace{11mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/
                 \big(\sqrt{\widehat{v_t}} + \epsilon \big)                                       \\
+            &\newline
+            &\hline \\[-1.ex]
             &\bf{return} \:  \theta_t                                                     \\[-1.ex]
-       \end{aligned}
+            &\newline
+            &\hline \\[-1.ex]
+       \end{array}
 
     .. warning::
         - This is an experimental optimizer API that is subject to change.
@@ -169,11 +177,10 @@ class AdamW(Optimizer):
         self.exp_avg_sq = self.parameters.clone(prefix="exp_avg_sq", init='zeros')
         self.state_step = Parameter(Tensor([-1], mstype.float32), "state_step")
         self.increase_tensor = Tensor(1, mstype.float32)
-        self.assignadd = P.AssignAdd()
         self.adamw_opt = gen.AdamW()
 
     def construct(self, gradients):
-        self.assignadd(self.state_step, self.increase_tensor)
+        self.state_step.add_(self.increase_tensor)
         for group_id, group in enumerate(self.param_groups):
             beta1, beta2 = group['betas']
             maximize = group.get("maximize")