mindspore 2.3.0__cp310-cp310-win_amd64.whl → 2.4.1__cp310-cp310-win_amd64.whl

mindspore/mint/nn/layer/normalization.py

@@ -0,0 +1,477 @@
+# Copyright 2020-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.
+# ============================================================================
+"""normalization for mint"""
+from __future__ import absolute_import
+from __future__ import division
+
+from typing import Optional
+import numpy as np
+import mindspore as ms
+from mindspore import mint
+from mindspore import ops
+from mindspore import Tensor
+from mindspore.common.parameter import Parameter
+from mindspore.common.initializer import initializer
+from mindspore import _checkparam as validator
+from mindspore.common import dtype as mstype
+from mindspore.nn.cell import Cell
+from mindspore.nn.layer.normalization import LayerNormExt as LayerNorm
+from mindspore.ops import group_norm
+
+
+class _NormBase(Cell):
+    """Common base of _InstanceNorm and _BatchNorm"""
+
+    def __init__(self,
+                 num_features: int,
+                 eps: float = 1e-5,
+                 momentum: float = 0.1,
+                 affine: bool = True,
+                 track_running_stats: bool = True,
+                 dtype=None
+                 ) -> None:
+        super(_NormBase, self).__init__()
+        self.shape = ops.Shape()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        self.track_running_stats = track_running_stats
+        self.dtype = dtype if dtype is not None else mstype.float32
+        if self.affine:
+            self.weight = Parameter(
+                Tensor(np.empty(num_features), dtype=self.dtype), name="weight")
+            self.bias = Parameter(
+                Tensor(np.empty(num_features), dtype=self.dtype), name="bias")
+            self.weight: Optional[Parameter]
+            self.bias: Optional[Parameter]
+        else:
+            self.weight = None
+            self.bias = None
+        if self.track_running_stats:
+            self.running_mean = Parameter(Tensor(np.zeros(num_features), dtype=self.dtype),
+                                          requires_grad=False, name="running_mean")
+            self.running_var = Parameter(Tensor(np.ones(num_features), dtype=self.dtype),
+                                         requires_grad=False, name="running_var")
+            self.running_mean: Optional[Tensor]
+            self.running_var: Optional[Tensor]
+            self.num_batches_tracked = Parameter(Tensor(0, dtype=ms.float32),
+                                                 requires_grad=False, name="num_batches_tracked")
+            self.num_batches_tracked: Optional[Tensor]
+        else:
+            self.running_mean = None
+            self.running_var = None
+            self.num_batches_tracked = None
+        self.reset_parameters()
+
+    def reset_running_stats(self) -> None:
+        """init parameters"""
+
+        if self.track_running_stats:
+            zero_running_mean = Tensor(
+                np.zeros(self.num_features), dtype=self.dtype)
+            one_running_var = Tensor(
+                np.ones(self.num_features), dtype=self.dtype)
+            zero_num_batches_tracked = Tensor(0, dtype=ms.float32)
+
+            ops.assign(self.running_mean, zero_running_mean)
+            ops.assign(self.running_var, one_running_var)
+            ops.assign(self.num_batches_tracked, zero_num_batches_tracked)
+
+    def reset_parameters(self) -> None:
+        self.reset_running_stats()
+        if self.affine:
+            one_weight = Tensor(np.ones(self.num_features), dtype=self.dtype)
+            zero_bias = Tensor(np.zeros(self.num_features), dtype=self.dtype)
+
+            ops.assign(self.weight, one_weight)
+            ops.assign(self.bias, zero_bias)
+
+    def _check_input_dim(self, input):
+        raise NotImplementedError
+
+    def extend_repr(self):
+        return 'num_features={}, eps={}, momentum={}, affine={}, track_running_stats={}'.format(
+            self.num_features, self.eps, self.momentum, self.affine, self.track_running_stats)
+
+
+class _BatchNorm(_NormBase):
+    """common base of BatchNormXxx"""
+
+    def __init__(
+            self,
+            num_features: int,
+            eps=1e-5,
+            momentum=0.1,
+            affine=True,
+            track_running_stats=True,
+            dtype=None) -> None:
+        super(_BatchNorm, self).__init__(num_features, eps, momentum, affine, track_running_stats,
+                                         dtype)
+        self.training = True
+
+    def _check_input_dim(self, input):
+        raise NotImplementedError
+
+    def construct(self, input):
+        self._check_input_dim(input)
+
+        if self.momentum is None:
+            exponential_average_factor = 0.0
+        else:
+            exponential_average_factor = self.momentum
+
+        if self.training and self.track_running_stats:
+            if self.num_batches_tracked is not None:
+                num_batches_tracked_one = Tensor(1, dtype=ms.float32)
+                ops.assign_add(self.num_batches_tracked,
+                               num_batches_tracked_one)
+                if self.momentum is None:
+                    exponential_average_factor = float(1.0 / self.num_batches_tracked)
+                else:
+                    exponential_average_factor = self.momentum
+
+        if self.training:
+            bn_training = True
+        else:
+            bn_training = (self.running_mean is None) and (
+                self.running_var is None)
+
+        return mint.functional.batch_norm(
+            input,
+            self.running_mean
+            if not self.training or self.track_running_stats
+            else None,
+            self.running_var if not self.training or self.track_running_stats
+            else None,
+            self.weight,
+            self.bias,
+            bn_training,
+            exponential_average_factor,
+            self.eps,
+        )
+
+
+class BatchNorm1d(_BatchNorm):
+    r"""
+    Applies Batch Normalization over a 2D or 3D input as described in the paper
+    `Batch Normalization: Accelerating Deep Network Training by Reducing
+    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`_ .
+
+    .. math::
+
+        y = \frac{x - mean}{\sqrt{variance + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the number of features or channels of the input). By default, the
+    elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0.
+
+    .. warning::
+        This API does not support Dynamic Rank.
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        num_features (int): `C` from an expected input of shape :math:`(N, C, L)`.
+        eps (float, optional): a value added to the denominator for numerical stability.
+            Default: ``1e-5`` .
+        momentum (float, optional): the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average. Default: ``0.1`` .
+        affine (bool, optional): a boolean value that when set to ``True``, this cell has
+            learnable affine parameters. Default: ``True`` .
+        track_running_stats (bool, optional): a boolean value that when set to ``True``, this
+            cell tracks the running mean and variance, and when set to ``False``,
+            this cell does not track such statistics. And this cell always uses batch statistics
+            in both training and eval modes. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`, optional): Dtype of Parameters. Default: ``None`` .
+
+    Inputs:
+        - **input** (Tensor) - The input with shape :math:`(N, C)` or :math:`(N, C, L)`,
+          where :math:`N` means batch, :math:`C` means the number of feature or the number of channel,
+          and :math:`L` is the length of sequence.
+
+    Outputs:
+        Tensor, has the same type and shape as `input`.
+
+    Raises:
+        TypeError: If `num_features` is not a int number.
+        TypeError: If `eps` is not a float.
+        ValueError: If `num_features` is less than 1.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> input_x = mindspore.Tensor([[0.7, 0.5, 0.5, 0.6], [0.5, 0.4, 0.6, 0.9]])
+        >>> net = mint.nn.BatchNorm1d(4)
+        >>> output = net(input_x)
+        >>> print(output)
+        [[ 0.99950075 0.9980011 -0.9980068 -0.9997783]
+         [-0.9995012 -0.99799967 0.9980068  0.9997778]]
+    """
+
+    def _check_input_dim(self, input):
+        shape = self.shape(input)
+        dim = len(shape)
+        if dim != 2 and dim != 3:
+            raise ValueError(
+                "expected 2D or 3D input (got {}D input)".format(dim)
+            )
+
+
+class BatchNorm2d(_BatchNorm):
+    r"""
+    Applies Batch Normalization over a 4D input as described in the paper
+    `Batch Normalization: Accelerating Deep Network Training by Reducing
+    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`_ .
+
+    .. math::
+
+        y = \frac{x - mean}{\sqrt{variance + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the number of features or channels of the input). By default, the
+    elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0.
+
+    .. warning::
+        This API does not support Dynamic Rank.
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        num_features (int): `C` from an expected input of shape :math:`(N, C, H, W)`.
+        eps (float, optional): a value added to the denominator for numerical stability.
+            Default: ``1e-5`` .
+        momentum (float, optional): the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average. Default: ``0.1`` .
+        affine (bool, optional): a boolean value that when set to ``True``, this cell has
+            learnable affine parameters. Default: ``True`` .
+        track_running_stats (bool, optional): a boolean value that when set to ``True``, this
+            cell tracks the running mean and variance, and when set to ``False``,
+            this cell does not track such statistics. And this cell always uses batch statistics
+            in both training and eval modes. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`, optional): Dtype of Parameters. Default: ``None`` .
+
+    Inputs:
+        - **input** (Tensor) - The input with shape :math:`(N, C, H, W)`.
+
+    Outputs:
+        Tensor, has the same type and shape as `input`.
+
+    Raises:
+        TypeError: If `num_features` is not a int number.
+        TypeError: If `eps` is not a float.
+        ValueError: If `num_features` is less than 1.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> input_x = mindspore.Tensor([0.3, 0.4, 0.5, 0.3])
+        >>> input_x = input_x.reshape((2, 2, 1, 1))
+        >>> net = mint.nn.BatchNorm2d(2)
+        >>> output = net(input_x)
+        >>> print(output)
+        [[[[-0.99950075]]
+          [[0.9980087]]]
+          [[[0.999501]]
+          [[-0.9980097]]]]
+    """
+
+    def _check_input_dim(self, input):
+        shape = self.shape(input)
+        dim = len(shape)
+        if dim != 4:
+            raise ValueError(
+                "expected 4D input (got {}D input)".format(dim)
+            )
+
+
+class BatchNorm3d(_BatchNorm):
+    r"""
+    Applies Batch Normalization over a 5D input as described in the paper
+    `Batch Normalization: Accelerating Deep Network Training by Reducing
+    Internal Covariate Shift <https://arxiv.org/abs/1502.03167>`_ .
+
+    .. math::
+
+        y = \frac{x - mean}{\sqrt{variance + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and :math:`\gamma` and :math:`\beta` are learnable parameter vectors
+    of size `C` (where `C` is the number of features or channels of the input). By default, the
+    elements of :math:`\gamma` are set to 1 and the elements of :math:`\beta` are set to 0.
+
+    .. warning::
+        This API does not support Dynamic Rank.
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        num_features (int): `C` from an expected input of shape :math:`(N, C, D, H, W)`.
+        eps (float, optional): a value added to the denominator for numerical stability.
+            Default: ``1e-5`` .
+        momentum (float, optional): the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average. Default: ``0.1`` .
+        affine (bool, optional): a boolean value that when set to ``True``, this cell has
+            learnable affine parameters. Default: ``True`` .
+        track_running_stats (bool, optional): a boolean value that when set to ``True``, this
+            cell tracks the running mean and variance, and when set to ``False``,
+            this cell does not track such statistics. And this cell always uses batch statistics
+            in both training and eval modes. Default: ``True`` .
+        dtype (:class:`mindspore.dtype`, optional): Dtype of Parameters. Default: ``None`` .
+
+    Inputs:
+        - **input** (Tensor) - The input with shape :math:`(N, C, D, H, W)`.
+
+    Outputs:
+        Tensor, has the same type and shape as `input`.
+
+    Raises:
+        TypeError: If `num_features` is not a int number.
+        TypeError: If `eps` is not a float.
+        ValueError: If `num_features` is less than 1.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> input_x = mindspore.Tensor([0.1, 0.9, 1.2, 2.3])
+        >>> input_x = input_x.reshape((1, 2, 1, 1, 2))
+        >>> net = mint.nn.BatchNorm3d(2)
+        >>> output = net(input_x)
+        >>> print(output)
+        [[[[[-0.9999688 0.99996865]]]
+          [[[-0.9999833 06.9999831]]]]]
+    """
+
+    def _check_input_dim(self, input):
+        shape = self.shape(input)
+        dim = len(shape)
+        if dim != 5:
+            raise ValueError(
+                "expected 5D input (got {}D input)".format(dim)
+            )
+
+
+class GroupNorm(Cell):
+    r"""
+    Group Normalization over a mini-batch of inputs.
+
+    Group Normalization is widely used in recurrent neural networks. It applies
+    normalization on a mini-batch of inputs for each single training case as described
+    in the paper `Group Normalization <https://arxiv.org/pdf/1803.08494.pdf>`_.
+
+    Group Normalization divides the channels into groups and computes within each group
+    the mean and variance for normalization, and it performs very stable over a wide
+    range of batch size. :math:`\gamma` and :math:`\beta` are trainable scale and shift.
+    It can be described using the following formula:
+
+    .. math::
+        y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    where :math:`\gamma` is `weight`, :math:`\beta` is `bias`, and :math:`\epsilon` is `eps`.
+
+    Args:
+        num_groups (int): The number of groups to be divided along the channel dimension.
+        num_channels (int): The number of input channels.
+        eps (float, optional): A value added to the denominator for numerical stability. Default: ``1e-05`` .
+        affine (bool, optional): The parameters, such as :math:`\gamma` and :math:`\beta`, are learnable
+            when set to ``true`` . Default: ``True`` .
+        dtype (:class:`mindspore.dtype`, optional): Dtype of Parameters. Default: ``None`` .
+
+    Inputs:
+        - **input** (Tensor) - The input feature with shape :math:`(N, C, *)`, where :math:`*` means, any number of
+          additional dimensions.
+
+    Outputs:
+        Tensor, the normalized and scaled offset tensor, has the same shape and data type as the `x`.
+
+    Raises:
+        TypeError: If `num_groups` or `num_channels` is not an int.
+        TypeError: If `eps` is not a float.
+        TypeError: If `affine` is not a bool.
+        ValueError: If `num_groups` or `num_channels` is less than 1.
+        ValueError: If `num_channels` is not divided by `num_groups`.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> group_norm_op = ms.mint.nn.GroupNorm(2, 2)
+        >>> x = ms.Tensor(np.ones([1, 2, 4, 4], np.float32))
+        >>> output = group_norm_op(x)
+        >>> print(output)
+        [[[[0. 0. 0. 0.]
+           [0. 0. 0. 0.]
+           [0. 0. 0. 0.]
+           [0. 0. 0. 0.]]
+          [[0. 0. 0. 0.]
+           [0. 0. 0. 0.]
+           [0. 0. 0. 0.]
+           [0. 0. 0. 0.]]]]
+    """
+
+    def __init__(self, num_groups, num_channels, eps=1e-05, affine=True, dtype=None):
+        """Initialize GroupNorm."""
+        super(GroupNorm, self).__init__()
+        ms_dtype = mstype.float32 if dtype is None else dtype
+        gamma_init = 'ones'
+        beta_init = 'zeros'
+
+        self.num_groups = validator.check_positive_int(
+            num_groups, "num_groups", self.cls_name)
+        self.num_channels = validator.check_positive_int(
+            num_channels, "num_channels", self.cls_name)
+        if num_channels % num_groups != 0:
+            raise ValueError(f"For '{self.cls_name}', the 'num_channels' must be divided by 'num_groups', "
+                             f"but got 'num_channels': {num_channels}, 'num_groups': {num_groups}.")
+        self.eps = validator.check_value_type(
+            'eps', eps, (float,), type(self).__name__)
+        self.affine = validator.check_bool(
+            affine, arg_name="affine", prim_name=self.cls_name)
+
+        self.gamma = Parameter(initializer(
+            gamma_init, self.num_channels, dtype=ms_dtype), name="gamma", requires_grad=affine)
+        self.beta = Parameter(initializer(
+            beta_init, self.num_channels, dtype=ms_dtype), name="beta", requires_grad=affine)
+
+    def _cal_output(self, x):
+        """calculate groupnorm output"""
+        return group_norm(x, self.num_groups, self.gamma, self.beta, self.eps)
+
+    def extend_repr(self):
+        return 'num_groups={}, num_channels={}, eps={}, affine={}'.format(
+            self.num_groups, self.num_channels, self.eps, self.affine)
+
+    def construct(self, input):
+        output = self._cal_output(input)
+        return output
+
+
+__all__ = [
+    'GroupNorm',
+    'BatchNorm1d',
+    'BatchNorm2d',
+    'BatchNorm3d',
+    'LayerNorm',
+]

mindspore/mint/nn/layer/pooling.py

@@ -0,0 +1,110 @@
+# Copyright 2020-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.
+# ============================================================================
+"""normalization for mint"""
+from __future__ import absolute_import
+from __future__ import division
+
+from mindspore import mint
+from mindspore.nn.cell import Cell
+
+
+class _AdaptiveAvgPoolNd(Cell):
+    """Common base of AdaptiveAvgPoolNd"""
+
+    def __init__(self, output_size) -> None:
+        super(_AdaptiveAvgPoolNd, self).__init__()
+        self.output_size = output_size
+
+    def extend_repr(self):
+        return 'output_size={}'.format(self.output_size)
+
+
+class AdaptiveAvgPool1d(_AdaptiveAvgPoolNd):
+    r"""
+    Applies a 1D adaptive average 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 (int): the target output size :math:`L_{out}` .
+
+    Inputs:
+        - **input** (Tensor) - The input with shape :math:`(N, C, L_{in})` or :math:`(C, L_{in})` .
+
+    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.AdaptiveAvgPool1d(3)
+        >>> output = net(input)
+        >>> print(output)
+        [[[2. 1. 2.]
+          [2. 3. 5.]]]
+    """
+
+    def construct(self, input):
+        return mint.nn.functional.adaptive_avg_pool1d(input, self.output_size)
+
+
+class AdaptiveAvgPool2d(_AdaptiveAvgPoolNd):
+    r"""
+    Applies a 2D adaptive average pooling over an input signal composed of several input planes.
+
+    The output is of size :math:`H x W` , 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[int])): the target output size of the image of the form :math:`H x W` .
+            Can be a tuple :math:`(H, W)` or a single :math:`H` for square image :math:`H x H` .
+            :math:`H` and :math:`W` can be either a ``int`` , or ``None`` which means the size will
+            be the same as that of the input.
+
+    Inputs:
+        - **input** (Tensor) - The input with shape :math:`(N, C, H, W)` or :math:`(C, H, W)` .
+
+    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.AdaptiveAvgPool2d((2, 2))
+        >>> output = net(input)
+        >>> print(output)
+        [[[1.5 1.5]
+          [2.5 4. ]]]
+    """
+
+    def construct(self, input):
+        return mint.nn.functional.adaptive_avg_pool2d(input, self.output_size)
+
+
+__all__ = [
+    'AdaptiveAvgPool2d',
+    'AdaptiveAvgPool1d',
+]

mindspore/mint/optim/adamw.py

@@ -23,18 +23,24 @@ from mindspore.ops import auto_generate as gen
 from mindspore.experimental.optim.optimizer import Optimizer
 from mindspore import _checkparam as validator
 
-_optim_adamw_opt = C.MultitypeFuncGraph("optim_adamw_opt")
 hyper_map = C.HyperMap()
 
 
-@_optim_adamw_opt.register("Function", "Float", "Float", "Float", "Float", "Float", "Tensor", "Bool", "Bool", "Tensor",
-                           "Tensor", "Tensor", "Tensor", "Tensor")
+def _run_optim_adamw_amsgrad_opt(opt, beta1, beta2, lr, eps, weight_decay, step, amsgrad, maximize, parameters, grads,
+                                 exp_avg, exp_avg_sq, max_exp_avg_sq):
+    """Apply adamw optimizer to the weight parameter."""
+    success = True
+    opt(parameters, exp_avg, exp_avg_sq, max_exp_avg_sq, grads, step, lr, beta1, beta2, weight_decay, eps, amsgrad,
+        maximize)
+    return success
+
+
 def _run_optim_adamw_opt(opt, beta1, beta2, lr, eps, weight_decay, step, amsgrad, maximize, parameters, grads, exp_avg,
-                         exp_avg_sq, max_exp_avg_sq):
+                         exp_avg_sq):
     """Apply adamw optimizer to the weight parameter."""
     success = True
-    opt(parameters, exp_avg, exp_avg_sq, max_exp_avg_sq, P.Cast()(grads, F.dtype(parameters)), step, lr, beta1, beta2,
-        weight_decay, eps, amsgrad, maximize)
+    opt(parameters, exp_avg, exp_avg_sq, exp_avg_sq, grads, step, lr, beta1, beta2, weight_decay, eps, amsgrad,
+        maximize)
     return success
 
 
@@ -156,15 +162,14 @@ class AdamW(Optimizer):
         defaults = dict(lr=lr, betas=betas, eps=eps,
                         weight_decay=weight_decay, amsgrad=amsgrad,
                         maximize=maximize)
+        self.max_v_group = True
         super(AdamW, 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.max_exp_avg_sq = self.parameters.clone(prefix="max_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.op_cast = P.Cast()
         self.adamw_opt = gen.AdamW()
 
     def construct(self, gradients):
@@ -177,9 +182,17 @@ class AdamW(Optimizer):
             lr = group.get("lr")
             grads = tuple(gradients[start_id: end_id])
 
-            self.hyper_map(F.partial(_optim_adamw_opt, self.adamw_opt, beta1, beta2, float(lr),
-                                     group.get("eps"), group.get("weight_decay"), self.state_step,
-                                     group.get("amsgrad"), maximize),
-                           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])
+            if group.get("amsgrad"):
+                self.hyper_map(F.partial(_run_optim_adamw_amsgrad_opt, self.adamw_opt, beta1, beta2, float(lr),
+                                         group.get("eps"), group.get("weight_decay"), self.state_step,
+                                         group.get("amsgrad"), maximize),
+                               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(_run_optim_adamw_opt, self.adamw_opt, beta1, beta2, float(lr),
+                                         group.get("eps"), group.get("weight_decay"), self.state_step,
+                                         group.get("amsgrad"), maximize),
+                               self.parameters[start_id: end_id], grads, self.exp_avg[start_id: end_id],
+                               self.exp_avg_sq[start_id: end_id])
+
         return True