mindspore 2.2.14__cp39-cp39-win_amd64.whl → 2.3.0__cp39-cp39-win_amd64.whl

mindspore/nn/dynamic_lr.py

@@ -27,7 +27,7 @@ def piecewise_constant_lr(milestone, learning_rates):
     Calculate learning rate by the given `milestone` and `learning_rates`. Let the value of `milestone` be
     :math:`(M_1, M_2, ..., M_t, ..., M_N)` and the value of `learning_rates` be :math:`(x_1, x_2, ..., x_t, ..., x_N)`.
     N is the length of `milestone`.
-    Let the output learning rate be :math:`y`, then for the :math:`i`-th step, the formula of
+    Let the output learning rate be :math:`y[i]`, then for the :math:`i`-th step, the formula of
     computing decayed_learning_rate[i] is:
 
     .. math::
@@ -360,7 +360,7 @@ def polynomial_decay_lr(learning_rate, end_learning_rate, total_step, step_per_e
         step_per_epoch (int): The number of steps in per epoch.
         decay_epoch (int): Number of epochs to decay over.
         power (float): The power of polynomial. It must be greater than 0.
-        update_decay_epoch (bool): If true, update `decay_epoch`. Default: ``False`` .
+        update_decay_epoch (bool): If ``true``, update `decay_epoch`. Default: ``False`` .
 
     Returns:
         list[float]. The size of list is `total_step`.

mindspore/nn/extend/__init__.py

@@ -0,0 +1,29 @@
+# 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.
+# ============================================================================
+"""
+nn Extend.
+"""
+from __future__ import absolute_import
+
+from mindspore.nn.extend.embedding import Embedding
+from mindspore.nn.extend.basic import Linear
+from mindspore.nn.extend.pooling import MaxPool2d
+from mindspore.nn.extend import layer
+from mindspore.nn.extend.layer import *
+
+__all__ = ['Embedding', 'Linear', 'MaxPool2d']
+__all__.extend(layer.__all__)
+
+__all__.sort()

mindspore/nn/extend/basic.py

@@ -0,0 +1,140 @@
+# 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.
+# ============================================================================
+
+"""basic"""
+from __future__ import absolute_import
+
+import math
+
+import mindspore.common.dtype as mstype
+from mindspore import _checkparam as Validator
+from mindspore._extends import cell_attr_register
+from mindspore.common.initializer import initializer, HeUniform, Uniform
+from mindspore.common.parameter import Parameter
+from mindspore.common.tensor import Tensor
+from mindspore.nn.cell import Cell
+from mindspore.ops import operations as P
+
+__all__ = ['Linear']
+
+
+class Linear(Cell):
+    r"""
+    The linear connected layer.
+
+    Applies linear connected layer for the input. This layer implements the operation as:
+
+    .. math::
+        \text{outputs} = X * kernel + bias
+
+    where :math:`X` is the input tensors, :math:`\text{kernel}` is a weight matrix with the same
+    data type as the :math:`X` created by the layer, and :math:`\text{bias}` is a bias vector
+    with the same data type as the :math:`X` created by the layer (only if has_bias is True).
+
+    Args:
+        in_features (int): The number of features in the input space.
+        out_features (int): The number of features in the output space.
+        bias (bool): Specifies whether the layer uses a bias vector :math:`\text{bias}`. Default: ``True``.
+        weight_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable weight_init parameter. The dtype
+            is same as `x`. The values of str refer to the function `initializer`. Default: ``None`` ,
+            weight will be initialized using HeUniform.
+        bias_init (Union[Tensor, str, Initializer, numbers.Number]): The trainable bias_init parameter. The dtype is
+            same as `x`. The values of str refer to the function `initializer`. Default: ``None`` ,
+            bias will be initialized using Uniform.
+        dtype (:class:`mindspore.dtype`): Data type of Parameter. Default: ``None`` .
+
+    Inputs:
+        - **x** (Tensor) - Tensor of shape :math:`(*, in\_features)`. The `in_features` in `Args` should be equal
+          to :math:`in\_features` in `Inputs`.
+
+    Outputs:
+        Tensor of shape :math:`(*, out\_features)`.
+
+    Raises:
+        TypeError: If `in_features` or `out_features` is not an int.
+        TypeError: If `bias` is not a bool.
+        ValueError: If length of shape of `weight_init` is not equal to 2 or shape[0] of `weight_init`
+                    is not equal to `out_features` or shape[1] of `weight_init` is not equal to `in_features`.
+        ValueError: If length of shape of `bias_init` is not equal to 1
+                    or shape[0] of `bias_init` is not equal to `out_features`.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor
+        >>> from mindspore import nn
+        >>> import numpy as np
+        >>> x = Tensor(np.array([[180, 234, 154], [244, 48, 247]]), mindspore.float32)
+        >>> net = nn.extend.Linear(3, 4)
+        >>> output = net(x)
+        >>> print(output.shape)
+        (2, 4)
+    """
+
+    @cell_attr_register(attrs=['has_bias'])
+    def __init__(self,
+                 in_features,
+                 out_features,
+                 bias=True,
+                 weight_init=None,
+                 bias_init=None,
+                 dtype=None):
+        """Initialize Linear."""
+        super(Linear, self).__init__()
+        self.in_features = Validator.check_positive_int(
+            in_features, "in_features", self.cls_name)
+        self.out_features = Validator.check_positive_int(
+            out_features, "out_features", self.cls_name)
+        self.has_bias = Validator.check_bool(
+            bias, "has_bias", self.cls_name)
+        self.dense = P.Dense()
+        if dtype is None:
+            dtype = mstype.float32
+        if isinstance(weight_init, Tensor):
+            if weight_init.ndim != 2 or weight_init.shape[0] != out_features or \
+                    weight_init.shape[1] != in_features:
+                raise ValueError(f"For '{self.cls_name}', weight init shape error. The ndim of 'weight_init' must "
+                                 f"be equal to 2, and the first dim must be equal to 'out_features', and the "
+                                 f"second dim must be equal to 'in_features'. But got 'weight_init': {weight_init}, "
+                                 f"'out_features': {out_features}, 'in_features': {in_features}.")
+        if weight_init is None:
+            weight_init = HeUniform(math.sqrt(5))
+        self.weight = Parameter(initializer(
+            weight_init, [out_features, in_features], dtype=dtype), name="weight")
+
+        self.bias = None
+        if self.has_bias:
+            if isinstance(bias_init, Tensor):
+                if bias_init.ndim != 1 or bias_init.shape[0] != out_features:
+                    raise ValueError(f"For '{self.cls_name}', bias init shape error. The ndim of 'bias_init' must "
+                                     f"be equal to 1, and the first dim must be equal to 'out_features'. But got "
+                                     f"'bias_init': {bias_init}, 'out_features': {out_features}.")
+            if bias_init is None:
+                bound = 1 / math.sqrt(in_features)
+                bias_init = Uniform(scale=bound)
+            self.bias = Parameter(initializer(
+                bias_init, [out_features], dtype=dtype), name="bias")
+
+    def construct(self, x):
+        x = self.dense(x, self.weight, self.bias)
+        return x
+
+    def extend_repr(self):
+        s = f'input_features={self.in_features}, output_features={self.out_features}'
+        if self.has_bias:
+            s += f', has_bias={self.has_bias}'
+        return s

mindspore/nn/extend/embedding.py

@@ -0,0 +1,143 @@
+# 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.
+# ============================================================================
+"""embedding"""
+from __future__ import absolute_import
+
+import mindspore.common.dtype as mstype
+from mindspore.common.initializer import Normal
+from mindspore import _checkparam as Validator
+from mindspore.nn.cell import Cell
+from mindspore import ops
+from mindspore.common.parameter import Parameter
+from mindspore.common.tensor import Tensor
+
+__all__ = ['Embedding']
+
+
+class Embedding(Cell):
+    r"""
+    Embedding layer.
+    Retrieve the word embeddings in weight stored in the layer using indices specified in `input`.
+
+    .. warning::
+        On Ascend, the behavior is unpredictable when the value of `input` is invalid.
+
+    Args:
+        num_embeddings (int): Size of the dictionary of embeddings.
+        embedding_dim (int): The size of each embedding vector.
+        padding_idx (int, optional): If the value is not None, the corresponding row of embedding vector
+            will not be updated in training. The value of embedding vector at `padding_idx` will default
+            to zeros when the Embedding layer is newly constructed. The value should be in range
+            `[-num_embeddings, num_embeddings)` if it's not ``None``. Default ``None``.
+        max_norm (float, optional): If the value is not None, firstly get the p-norm result of the embedding
+            vector specified by `input` where p is specified by `norm_type`; if the result is larger then `max_norm`,
+            update the embedding vector` with :math:`\frac{max\_norm}{result+1e^{-7}}`. Default ``None``.
+        norm_type (float, optional): Indicated the value of p in p-norm. Default ``2.0``.
+        scale_grad_by_freq (bool, optional): If ``True`` the gradients will be scaled by the inverse of frequency
+            of the index in `input`. Default ``False``.
+        _weight (Tensor, optional): Used to initialize the weight of Embedding. If ``None``, the weight will be
+            initialized from normal distribution :math:`{N}(\text{sigma=1.0}, \text{mean=0.0})`. Default ``None``.
+        dtype (mindspore.dtype, optional) : Dtype of Parameters. It is meaningless when `_weight` is not None.
+            Default: ``mindspore.float32``.
+
+    Inputs:
+        - **input** (Tensor) - The indices used to lookup in the embedding vector. The data type must be
+          mindspore.int32 or mindspore.int64, and the value should be in range `[0, num_embeddings)`.
+
+    Outputs:
+        Tensor, has the same data type as weight, the shape is :math:`(*input.shape, embedding\_dim)`.
+
+    Raises:
+        TypeError: If `num_embeddings` is not an int.
+        TypeError: If `embedding_dim` is not an int.
+        ValueError: If `padding_idx` is out of valid range.
+        TypeError: If `max_norm` is not a float.
+        TypeError: If `norm_type` is not a float.
+        TypeError: If `scale_grad_by_freq` is not a bool.
+        TypeError: If `dtype` is not one of mindspore.dtype.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> import numpy as np
+        >>> from mindspore import Tensor, nn
+        >>> input = Tensor([[1, 0, 1, 1], [0, 0, 1, 0]])
+        >>> embedding = nn.extend.Embedding(num_embeddings=10, embedding_dim=3)
+        >>> output = embedding(input)
+        >>> print(output)
+        [[[-0.0024154  -0.01203444  0.00811537]
+          [ 0.00233847 -0.00596091  0.00536799]
+          [-0.0024154  -0.01203444  0.00811537]
+          [-0.0024154  -0.01203444  0.00811537]]
+         [[ 0.00233847 -0.00596091  0.00536799]
+          [ 0.00233847 -0.00596091  0.00536799]
+          [-0.0024154  -0.01203444  0.00811537]
+          [ 0.00233847 -0.00596091  0.00536799]]]
+    """
+
+    def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=None, norm_type=2.0,
+                 scale_grad_by_freq=False, _weight=None, dtype=mstype.float32):
+        """Initialize Embedding."""
+        super().__init__()
+        self.num_embeddings = Validator.check_value_type(
+            'num_embeddings', num_embeddings, [int], self.cls_name)
+        self.embedding_dim = Validator.check_value_type(
+            'embedding_dim', embedding_dim, [int], self.cls_name)
+        Validator.check_subclass(
+            "dtype", dtype, mstype.number_type, self.cls_name)
+        self.dtype = dtype
+        self.padding_idx = padding_idx
+        if _weight is None:
+            init_tensor = Tensor(shape=[num_embeddings, embedding_dim], dtype=dtype, init=Normal(1, 0))
+            init_tensor = self._zero_weight_by_index(init_tensor)
+            self.weight = Parameter(init_tensor, name='weight')
+        else:
+            self.weight = Parameter(_weight)
+
+        self.max_norm = max_norm
+        if max_norm is not None:
+            self.max_norm = Validator.check_value_type('max_norm', max_norm, [float], self.cls_name)
+
+        self.norm_type = norm_type
+        if norm_type is not None:
+            self.norm_type = Validator.check_value_type('norm_type', norm_type,
+                                                        [float], self.cls_name)
+
+        self.scale_grad_by_freq = scale_grad_by_freq
+        if scale_grad_by_freq is not None:
+            self.scale_grad_by_freq = Validator.check_value_type('scale_grad_by_freq',
+                                                                 scale_grad_by_freq,
+                                                                 [bool], self.cls_name)
+
+    def _zero_weight_by_index(self, init_tensor):
+        if self.padding_idx is not None:
+            self.padding_idx = Validator.check_int_range(self.padding_idx, -self.num_embeddings, self.num_embeddings,
+                                                         Validator.INC_LEFT, "padding_idx", self.cls_name)
+            if isinstance(init_tensor, Tensor) and init_tensor.init is not None:
+                init_tensor = init_tensor.init_data()
+            init_tensor[self.padding_idx] = 0
+
+        return init_tensor
+
+    def construct(self, input):
+        return ops.embedding(input, self.weight, self.padding_idx, self.max_norm,
+                             self.norm_type, self.scale_grad_by_freq)
+
+    def extend_repr(self):
+        return f'num_embeddings={self.num_embeddings}, embedding_dim={self.embedding_dim}, '  \
+               f'padding_idx={self.padding_idx}, max_norm={self.max_norm}, norm_type={self.norm_type}, ' \
+               f'scale_grad_by_freq={self.scale_grad_by_freq}, dtype={self.dtype}'

mindspore/{rewrite/ast_creator_register.py → nn/extend/layer/__init__.py}

@@ -1,4 +1,4 @@
-# Copyright 2022 Huawei Technologies Co., Ltd
+# 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.
@@ -12,26 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
+"""
+Layer.
 
-"""Registry the relation."""
-
+The high-level components(Cells) used to construct the neural network.
+"""
 from __future__ import absolute_import
-from collections import UserDict
-
-
-class Registry(UserDict):
-    """Registry class for registry functions for creating ast node."""
-
-    def register(self, obj_str, obj):
-        """Register object by str."""
-        if isinstance(obj_str, str):
-            self[obj_str] = obj
 
-    def get(self, obj_str):
-        """Get the value by str."""
-        if not isinstance(obj_str, str):
-            raise TypeError("key for ast creator registry must be string.")
-        return self[obj_str]
+from mindspore.nn.extend.layer import normalization
+from mindspore.nn.extend.layer.normalization import *
 
+__all__ = []
 
-ast_creator_registry = Registry()
+__all__.extend(normalization.__all__)

mindspore/nn/extend/layer/normalization.py

@@ -0,0 +1,109 @@
+# 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"""
+from __future__ import absolute_import
+from __future__ import division
+
+from mindspore.ops import functional as F
+from mindspore.common.parameter import Parameter
+from mindspore.common.initializer import initializer
+from mindspore.common import dtype as mstype
+from mindspore.nn.cell import Cell
+
+__all__ = ['LayerNorm']
+
+
+class LayerNorm(Cell):
+    r"""
+    Applies Layer Normalization over a mini-batch of inputs.
+
+    Layer 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 `Layer Normalization <https://arxiv.org/pdf/1607.06450.pdf>`_. Unlike Batch
+    Normalization, Layer Normalization performs exactly the same computation at training and
+    testing time. It is applied across all channels and pixel but only one 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
+
+    Args:
+        normalized_shape (Union(tuple[int], list[int])): The normalized shape of `x` for LayerNorm
+        gamma_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\gamma` weight.
+            The values of str refer to the function `initializer` including ``'zeros'`` , ``'ones'`` ,
+            ``'xavier_uniform'`` , ``'he_uniform'`` , etc. Default: ``'ones'`` .
+        beta_init (Union[Tensor, str, Initializer, numbers.Number]): Initializer for the :math:`\beta` weight.
+            The values of str refer to the function `initializer` including ``'zeros'`` , ``'ones'`` ,
+            ``'xavier_uniform'`` , ``'he_uniform'`` , etc. Default: ``'zeros'`` .
+        eps (float): A value added to the denominator for numerical stability(:math:`\epsilon`). Default: ``1e-5`` .
+        elementwise_affine (bool): A bool value, When set to True, gamma and beta can be learned. Default: True.
+        dtype (:class:`mindspore.dtype`): Dtype of Parameters. Default: ``mstype.float32`` .
+
+    Inputs:
+        - **x** (Tensor) - The shape is :math:`(N, *)`, 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 `epsilon` is not a float.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> x = ms.Tensor(np.ones([20, 5, 10, 10]), ms.float32)
+        >>> shape1 = x.shape[1:]
+        >>> m = ms.nn.extend.LayerNorm(shape1)
+        >>> output = m(x).shape
+        >>> print(output)
+        (20, 5, 10, 10)
+    """
+
+    def __init__(self,
+                 normalized_shape,
+                 gamma_init='ones',
+                 beta_init='zeros',
+                 eps=1e-5,
+                 elementwise_affine=True,
+                 dtype=mstype.float32
+                 ):
+        """Initialize LayerNorm."""
+        super(LayerNorm, self).__init__()
+        if not isinstance(normalized_shape, (tuple, list)):
+            raise TypeError(f"For '{self.cls_name}', the type of 'normalized_shape' must be tuple[int] or list[int], "
+                            f"but got {normalized_shape} and the type is {type(normalized_shape)}.")
+        if not normalized_shape:
+            raise ValueError(
+                f"Expected normalized_shape to be at least 1-dimensional, i.e., containing at "
+                f"least one element, but got normalized_shape = {normalized_shape}"
+            )
+        self.normalized_shape = normalized_shape
+        self.epsilon = eps
+        self.gamma = Parameter(initializer(
+            gamma_init, normalized_shape, dtype=dtype), name="gamma", requires_grad=elementwise_affine)
+        self.beta = Parameter(initializer(
+            beta_init, normalized_shape, dtype=dtype), name="beta", requires_grad=elementwise_affine)
+
+    def construct(self, input_x):
+        y = F.layer_norm(input_x, self.normalized_shape, self.gamma.astype(input_x.dtype),
+                         self.beta.astype(input_x.dtype), self.epsilon)
+        return y
+
+    def extend_repr(self):
+        return 'normalized_shape={}, gamma{}, beta={}'.format(self.normalized_shape, self.gamma, self.beta)

mindspore/nn/extend/pooling.py

@@ -0,0 +1,117 @@
+#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.
+# ============================================================================
+"""pooling"""
+from __future__ import absolute_import
+
+from mindspore.ops.auto_generate.gen_ops_prim import MaxPoolWithIndices, MaxPoolWithMask
+from mindspore.nn.cell import Cell
+
+__all__ = ['MaxPool2d']
+
+
+class MaxPool2d(Cell):
+    r"""
+    Applies a 2D max pooling over an input Tensor which can be regarded as a composition of 2D planes.
+
+    Typically the input is of shape :math:`(N_{in}, C_{in}, H_{in}, W_{in})`, MaxPool2d outputs
+    regional maximum in the :math:`(H_{in}, W_{in})`-dimension. Given kernel size
+    :math:`(h_{ker}, w_{ker})` and stride :math:`(s_0, s_1)`, the operation is as follows.
+
+    .. math::
+        \text{output}(N_i, C_j, h, w) = \max_{m=0, \ldots, h_{ker}-1} \max_{n=0, \ldots, w_{ker}-1}
+        \text{input}(N_i, C_j, s_0 \times h + m, s_1 \times w + n)
+
+    .. warning::
+        Only support on Atlas training series.
+
+    Args:
+        kernel_size (Union[int, tuple[int]]): The size of kernel used to take the max value,
+            is an int number or a single element tuple that represents height and width are both kernel_size,
+            or a tuple of two int numbers that represent height and width respectively.
+            Default: ``1`` .
+        stride (Union[int, tuple[int], None]): The distance of kernel moving, an int number or a single element tuple
+            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], list[int])): Specifies the padding value of the pooling operation.
+            Default: ``0`` . `padding` can only be an integer or a tuple/list containing one or two integers. If
+            `padding` is an integer or a tuple/list containing one integer, it will be padded `padding` times in the
+            four directions of the input. If `padding` is a tuple/list containing two integers, it will be padded
+            `padding[0]` times in the up-down direction of the input and `padding[1]` times in the left-right direction
+            of the input.
+        dilation (Union(int, tuple[int])): The spacing between the elements of the kernel in convolution,
+            used to increase the receptive field of the pooling operation. If it is a tuple, it must contain one or two
+            integers. Default: ``1`` .
+        return_indices (bool): If ``True`` , the function will return both the result of max pooling and the indices of
+            the max elements. Default: ``False`` .
+        ceil_mode (bool): If ``True`` , use ceil to compute the output shape instead of floor. Default: ``False`` .
+
+    Inputs:
+        - **input** (Tensor) - Tensor of shape :math:`(N, C_{in}, H_{in}, W_{in})`.
+
+    Outputs:
+        If `return_indices` is ``False`` , return a Tensor `output`, else return a tuple (`output`, `argmax`).
+
+        - **output** (Tensor) - Maxpooling result, with shape :math:`(N_{out}, C_{out}, H_{out}, W_{out})`. It has the
+          same data type as `input`.
+        - **argmax** (Tensor) - Index corresponding to the maximum value. Data type is int32.
+
+        .. math::
+            H_{out} = \left\lfloor\frac{H_{in} + 2 * \text{padding[0]} - \text{dilation[0]}
+                \times (\text{kernel_size[0]} - 1) - 1}{\text{stride[0]}} + 1\right\rfloor
+
+        .. math::
+            W_{out} = \left\lfloor\frac{W_{in} + 2 * \text{padding[1]} - \text{dilation[1]}
+                \times (\text{kernel_size[1]} - 1) - 1}{\text{stride[1]}} + 1\right\rfloor
+
+    Raises:
+        TypeError: If `input` is not a Tensor.
+        ValueError: If length of shape of `input` is not equal to 4.
+        TypeError: If `kernel_size` , `stride` , `padding` or `dilation` is not int or tuple.
+        ValueError: If `kernel_size`, `stride` or `dilation` is less than 1.
+        ValueError: If `dilation` is not all 1.
+        ValueError: If `padding` is less than 0.
+        ValueError: If `padding` is more than half of `kernel_size`.
+        TypeError: If `ceil_mode` is not bool.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> pool = ms.nn.extend.MaxPool2d(kernel_size=3, stride=1)
+        >>> input = ms.Tensor(np.random.randint(0, 10, [1, 2, 4, 4]), ms.float32)
+        >>> output = pool(input)
+        >>> print(output.shape)
+        (1, 2, 2, 2)
+    """
+
+    def __init__(self, kernel_size=1, stride=None, padding=0, dilation=1, return_indices=False,
+                 ceil_mode=False):
+        """Initialize MaxPool2d."""
+        super(MaxPool2d, self).__init__()
+        self.return_indices = return_indices
+        strides = stride if (stride is not None) else kernel_size
+        if return_indices:
+            self.max_pool_func_ = MaxPoolWithIndices(kernel_size, strides, padding, dilation, ceil_mode)
+        else:
+            self.max_pool_func_ = MaxPoolWithMask(kernel_size, strides, padding, dilation, ceil_mode)
+
+    def construct(self, input):
+        out, indices = self.max_pool_func_(input)
+        if self.return_indices:
+            return out, indices
+        return out