mindspore 2.4.1__cp311-cp311-manylinux1_x86_64.whl → 2.5.0__cp311-cp311-manylinux1_x86_64.whl

mindspore/mint/nn/layer/_functions.py

@@ -0,0 +1,330 @@
+import mindspore
+from mindspore import Tensor
+from mindspore import context
+from mindspore.nn.cell import Cell
+from mindspore.ops.auto_generate.gen_ops_prim import BatchNormReduceGrad
+from mindspore.ops.auto_generate.gen_ops_prim import BatchNormElemtGrad
+from mindspore.communication import GlobalComm
+from mindspore.ops import ReduceOp
+from mindspore._c_expression import Tensor as Tensor_
+from mindspore.communication._comm_helper import _get_size_helper, HCCL_WORLD_COMM_GROUP
+from mindspore.ops._primitive_cache import _get_cache_prim
+from mindspore.ops import operations as P
+from mindspore import ops, mint
+
+
+DEFAULT_WORLD_COMM_GROUP = HCCL_WORLD_COMM_GROUP
+
+batch_norm_reduce_grad = BatchNormReduceGrad()
+batch_norm_elemt_grad = BatchNormElemtGrad()
+shape = P.Shape()
+
+
+def _deal_comm_outputs(output, async_op):
+    if isinstance(output, tuple):
+        if not async_op:
+            output[1].wait()
+            return output[0]
+        return output
+
+    if not async_op:
+        return output
+    return output
+
+
+def get_group_size(group=GlobalComm.WORLD_COMM_GROUP):
+    if not isinstance(group, str):
+        raise TypeError("For 'get_group_size', the argument 'group' must be type of string, "
+                        "but got 'group' type : {}.".format(type(group)))
+    return _get_size_helper(group=_get_group(group))
+
+
+def _contiguous(tensor):
+    if not tensor.is_contiguous() or tensor.storage_offset() != 0:
+        tensor = tensor.contiguous()
+    return tensor
+
+
+def _get_group(group):
+    """Return the world communication group if the `group` is `DEFAULT_WORLD_COMM_GROUP`."""
+    if group == DEFAULT_WORLD_COMM_GROUP:
+        return GlobalComm.WORLD_COMM_GROUP
+    return group
+
+
+def all_gather_into_tensor(tensor, group=GlobalComm.WORLD_COMM_GROUP, async_op=False):
+    if not isinstance(tensor, (Tensor, Tensor_)):
+        raise TypeError(
+            "For all_gather_into_tensor, the input tensor must be tensor")
+    group = _get_group(group)
+    tensor = _contiguous(tensor)
+    all_gather_op = _get_cache_prim(P.AllGather)(group=group)
+    output = all_gather_op(tensor)
+    return _deal_comm_outputs(output, async_op)
+
+
+def all_reduce(tensor, op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP, async_op=False):
+    if not isinstance(tensor, (Tensor, Tensor_)):
+        raise TypeError("For all_reduce, the input tensor must be tensor")
+    if not isinstance(op, str):
+        raise TypeError("For all_reduce, the input op type must be str")
+    if op not in ('sum', 'prod', 'min', 'max'):
+        raise TypeError(
+            "For all_reduce, the input op value must be one of sum, prod, min, max")
+    group = _get_group(group)
+    tensor = _contiguous(tensor)
+    all_reduce_op = _get_cache_prim(P.AllReduce)(op=op, group=group)
+    output = all_reduce_op(tensor)
+    return _deal_comm_outputs(output, async_op)
+
+
+def bprop_pynative(input_x, weight, bias, running_mean, running_var, eps, momentum,
+                   process_group, world_size, output, doutput):
+    _, mean_param, invstd_param, count_all_param = output
+    dout, _, _, _ = doutput
+
+    # 不支持 KBK模式
+    if not dout.is_contiguous():
+        dout = dout.contiguous()
+
+    grad_input = grad_weight = grad_bias = None
+
+    inputG = True
+    weightG = True
+    biasG = True
+
+    # calculate local stats as well as grad_weight / grad_bias
+    sum_dy, sum_dy_xmu, grad_weight, grad_bias = batch_norm_reduce_grad(
+        dout,
+        input_x,
+        mean_param,
+        invstd_param,
+        weight,
+        inputG,
+        weightG,
+        biasG
+    )
+
+    if inputG:
+        # synchronizing stats used to calculate input gradient.
+        sum_dy_shape = shape(sum_dy)
+        num_channels = sum_dy_shape[0]
+        combined = mint.cat([sum_dy, sum_dy_xmu], dim=0)
+
+        new_combined = all_reduce(combined, group=process_group)
+
+        sum_dy, sum_dy_xmu = mint.split(new_combined, num_channels)
+
+        # backward pass for gradient calculation
+        grad_input = batch_norm_elemt_grad(
+            dout,
+            input_x,
+            mean_param,
+            invstd_param,
+            weight,
+            sum_dy,
+            sum_dy_xmu,
+            count_all_param
+        )
+
+    # synchronizing of grad_weight / grad_bias is not needed as distributed
+    # training would handle all reduce.
+    if weight is None or not weightG:
+        grad_weight = None
+
+    if weight is None or not biasG:
+        grad_bias = None
+
+    return grad_input, grad_weight, grad_bias, None, None, None, None, None, None
+
+
+def bprop_kbk(input_x, weight, bias, running_mean, running_var, eps, momentum,
+              process_group, world_size, output, doutput):
+    _, mean_param, invstd_param, count_all_param = output
+    dout, _, _, _ = doutput
+
+    dout = dout.contiguous()
+
+    grad_input = grad_weight = grad_bias = None
+
+    inputG = True
+    weightG = True
+    biasG = True
+
+    # calculate local stats as well as grad_weight / grad_bias
+    sum_dy, sum_dy_xmu, grad_weight, grad_bias = batch_norm_reduce_grad(
+        dout,
+        input_x,
+        mean_param,
+        invstd_param,
+        weight,
+        inputG,
+        weightG,
+        biasG
+    )
+
+    if inputG:
+        # synchronizing stats used to calculate input gradient.
+        sum_dy_shape = shape(sum_dy)
+        num_channels = sum_dy_shape[0]
+        combined = mint.cat([sum_dy, sum_dy_xmu], dim=0)
+
+        new_combined = all_reduce(combined, group=process_group)
+
+        sum_dy, sum_dy_xmu = mint.split(new_combined, num_channels)
+
+        # backward pass for gradient calculation
+        grad_input = batch_norm_elemt_grad(
+            dout,
+            input_x,
+            mean_param,
+            invstd_param,
+            weight,
+            sum_dy,
+            sum_dy_xmu,
+            count_all_param
+        )
+
+    # synchronizing of grad_weight / grad_bias is not needed as distributed
+    # training would handle all reduce.
+    if weight is None or not weightG:
+        grad_weight = None
+
+    if weight is None or not biasG:
+        grad_bias = None
+
+    return grad_input, grad_weight, grad_bias, None, None, None, None, None, None
+
+
+def construct_pynative(input, weight, bias, running_mean, running_var, eps, momentum, process_group,
+                       world_size, self_num_features, self_world_size):
+    if self_world_size != world_size:
+        raise ValueError('World Size Error')
+    if not input.is_contiguous():
+        input = input.contiguous()
+    if weight is not None:
+        weight = weight.contiguous()
+
+    input_shape = shape(input)
+    input_numel = ops.numel(input)
+    size = int(input_numel // input_shape[1])
+    if size == 1 and world_size < 2:
+        raise ValueError(
+            'Expected more than 1 value per channel when training, got input size {}'.format(size))
+
+    # calculate mean/invstd for input.
+    mean, invstd = mint.batch_norm_stats(input, eps)
+    count = mint.full((1,), input_numel //
+                      input_shape[1], dtype=mean.dtype)
+
+    num_channels = input_shape[1]
+    if self_num_features != num_channels:
+        raise ValueError('Features Error')
+    # C, C, 1 -> (2C + 1)
+    combined = mint.cat([mean, invstd, count], dim=0)
+    # Use allgather instead of allreduce because count could be different across
+    # ranks, simple all reduce op can not give correct results.
+    # batch_norm_gather_stats_with_counts calculates global mean & invstd based on
+    # all gathered mean, invstd and count.
+    # world_size * (2C + 1)
+    combined = all_gather_into_tensor(combined, process_group)
+    combined = ops.reshape(combined, [world_size, -1])
+    # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
+    mean_val_all, invstd_val_all, count_val_all = mint.split(
+        combined, num_channels, dim=1)
+    # calculate global mean & invstd
+    mean, invstd = mint.batch_norm_gather_stats_with_counts(input, mean_val_all, invstd_val_all, running_mean,
+                                                            running_var, momentum, eps, count_val_all.view(-1))
+
+    # apply element-wise normalization
+    out = mint.batch_norm_elemt(input, weight, bias, mean, invstd, eps)
+    return (out, mean, invstd, count_val_all.view(-1))
+
+
+def construct_kbk(input, weight, bias, running_mean, running_var, eps, momentum, process_group,
+                  world_size, self_num_features, self_world_size):
+    if self_world_size != world_size:
+        raise ValueError('World Size Error')
+    input = input.contiguous()
+    if weight is not None:
+        weight = weight.contiguous()
+
+    input_shape = shape(input)
+    input_numel = ops.numel(input)
+    size = int(input_numel // input_shape[1])
+    if size == 1 and world_size < 2:
+        raise ValueError(
+            'Expected more than 1 value per channel when training, got input size {}'.format(size))
+
+    # calculate mean/invstd for input.
+    mean, invstd = mint.batch_norm_stats(input, eps)
+    count = mint.full((1,), input_numel //
+                      input_shape[1], dtype=mean.dtype)
+
+    num_channels = input_shape[1]
+    if self_num_features != num_channels:
+        raise ValueError('Features Error')
+    # C, C, 1 -> (2C + 1)
+    combined = mint.cat([mean, invstd, count], dim=0)
+    # Use allgather instead of allreduce because count could be different across
+    # ranks, simple all reduce op can not give correct results.
+    # batch_norm_gather_stats_with_counts calculates global mean & invstd based on
+    # all gathered mean, invstd and count.
+    # world_size * (2C + 1)
+    combined = all_gather_into_tensor(combined, process_group)
+    combined = ops.reshape(combined, [world_size, -1])
+    # world_size * (2C + 1) -> world_size * C, world_size * C, world_size * 1
+    mean_all, invstd_all, count_all = mint.split(
+        combined, num_channels, dim=1)
+    # calculate global mean & invstd
+    mean, invstd = mint.batch_norm_gather_stats_with_counts(
+        input,
+        mean_all,
+        invstd_all,
+        running_mean,
+        running_var,
+        momentum,
+        eps,
+        count_all.view(-1)
+    )
+
+    # apply element-wise normalization
+    out = mint.batch_norm_elemt(input, weight, bias, mean, invstd, eps)
+    return (out, mean, invstd, count_all.view(-1))
+
+
+class SyncBatchNormInner(Cell):
+    def __init__(self, self_num_features, self_world_size):
+        super(SyncBatchNormInner, self).__init__()
+        self.num_features = self_num_features
+        self.world_size = self_world_size
+        self.mode = context.get_context("mode")
+        if self.mode == 1:
+            self.fn_bprop = bprop_pynative
+            self.fn_construct = construct_pynative
+        else:
+            self.fn_bprop = bprop_kbk
+            self.fn_construct = construct_kbk
+
+    def construct(self, input, weight, bias, running_mean, running_var, eps, momentum, process_group, world_size):
+        return self.fn_construct(input, weight, bias, running_mean, running_var, eps, momentum, process_group,
+                                 world_size, self.num_features, self.world_size)
+
+    def bprop(self, input_x, weight, bias, running_mean, running_var, eps, momentum,
+              process_group, world_size, output, doutput):
+        return self.fn_bprop(input_x, weight, bias, running_mean, running_var, eps, momentum,
+                             process_group, world_size, output, doutput)
+
+
+class _SyncBatchNorm(Cell):
+    def __init__(self, num_features, world_size, dtype=mindspore.float32):
+        super(_SyncBatchNorm, self).__init__()
+        self.num_features = num_features
+        self.world_size = world_size
+        self.inner = SyncBatchNormInner(self.num_features, self.world_size)
+
+    def construct(self, input, weight, bias, running_mean, running_var, eps, momentum, process_group, world_size):
+        res = self.inner(input, weight, bias, running_mean,
+                         running_var, eps, momentum, process_group, world_size)
+        output, _, _, _ = res
+        return output

mindspore/mint/nn/layer/activation.py

@@ -84,7 +84,7 @@ class LogSigmoid(Cell):
     Logsigmoid is defined as:
 
     .. math::
-        \text{logsigmoid}(x_{i}) = \log(\frac{1}{1 + \exp(-x_i)}),
+        \text{LogSigmoid}(x_{i}) = \log(\frac{1}{1 + \exp(-x_i)}),
 
     where :math:`x_{i}` is the element of the input.
 
@@ -127,7 +127,175 @@ class LogSigmoid(Cell):
         return mint.nn.functional.logsigmoid(input)
 
 
+class ELU(Cell):
+    r"""
+    Exponential Linear Unit activation function
+
+    Applies the exponential linear unit function element-wise.The activation function is defined as:
+
+    .. math::
+        ELU_{i} =
+        \begin{cases}
+        x_i, &\text{if } x_i \geq 0; \cr
+        \alpha * (\exp(x_i) - 1), &\text{otherwise.}
+        \end{cases}
+
+    where :math:`x_i` represents the element of the input and :math:`\alpha` represents the `alpha` parameter.
+
+    ELU Activation Function Graph:
+
+    .. image:: ../images/ELU.png
+        :align: center
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        alpha (float, optional): The alpha value of ELU, the data type is float. Default: ``1.0`` .
+
+    Inputs:
+        - **input** (Tensor) - The input of ELU is a Tensor of any dimension.
+
+    Outputs:
+        Tensor, with the same type and shape as the `input`.
+
+    Raises:
+        TypeError: If `alpha` is not a float.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> import numpy as np
+        >>> input = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float32)
+        >>> elu = mint.nn.ELU()
+        >>> result = elu(input)
+        >>> print(result)
+        [-0.63212055  -0.86466473  0.  2.  1.]
+    """
+
+    def __init__(self, alpha=1.0):
+        """Initialize ELU."""
+        super(ELU, self).__init__()
+        self.alpha = alpha
+
+    def construct(self, input):
+        return mint.nn.functional.elu(input, self.alpha)
+
+
+class GLU(Cell):
+    r"""
+    Computes GLU (Gated Linear Unit activation function) of the input tensor.
+
+    .. math::
+        {GLU}(a, b)= a \otimes \sigma(b)
+
+    where :math:`a` is the first half of the `input` Tensor after `input` is split and :math:`b` is the second half.
+
+    Here :math:`\sigma` is the sigmoid function, and :math:`\otimes` is the Hadamard product.
+    See `Language Modeling with Gated Convluational Networks <https://arxiv.org/abs/1612.08083>`_ .
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        dim (int, optional): The dimension to split the input `input`. The value range is `[-r, r)` where `r`
+            is the number of dimensions of `input`. Default: ``-1`` , the last dimension in `input`.
+
+    Inputs:
+        - **input** (Tensor) - Tensor to be calculated. Dtype is floating point and the shape
+          is :math:`(\ast_1, N, \ast_2)` where `*` means, any number of additional dimensions. :math:`N`
+          is required to be an even number, where :math:`N` is the size of `input` on the dimension
+          selected by `dim`.
+
+    Outputs:
+        Tensor, the same dtype as the `input`, with the shape :math:`(\ast_1, M, \ast_2)` where :math:`M=N/2`.
+
+    Raises:
+        TypeError: If `input` is not a Tensor or `dim` is not an int.
+        IndexError: If the value of `dim` is out of the range of `[-r, r)`, where `r` is the number
+            of dimensions of `input`.
+        RuntimeError: If dtype of `input` is not supported.
+        RuntimeError: If the length of `input` in the dimension selected by `dim` is not even.
+
+    Supported Platforms:
+        ``Ascend`` ``CPU``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> m = ms.mint.nn.GLU()
+        >>> input = ms.Tensor([[0.1,0.2,0.3,0.4],[0.5,0.6,0.7,0.8]])
+        >>> output = m(input)
+        >>> print(output)
+        [[0.05744425 0.11973753]
+         [0.33409387 0.41398472]]
+    """
+
+    def __init__(self, dim=-1):
+        """Initialize GLU."""
+        super().__init__("GLU")
+        self.dim = dim
+
+    def construct(self, input):
+        return mint.nn.functional.glu(input, self.dim)
+
+
+class Tanh(Cell):
+    r"""
+    Applies the Tanh function element-wise, returns a new tensor with the hyperbolic tangent of the elements of input.
+
+    Tanh function is defined as:
+
+    .. math::
+        tanh(x_i) = \frac{\exp(x_i) - \exp(-x_i)}{\exp(x_i) + \exp(-x_i)} = \frac{\exp(2x_i) - 1}{\exp(2x_i) + 1},
+
+    where :math:`x_i` is an element of the input Tensor.
+
+    Tanh Activation Function Graph:
+
+    .. image:: ../images/Tanh.png
+        :align: center
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Inputs:
+        - **input** (Tensor) - Tensor of any dimension, input with data type of float16 or float32.
+
+    Outputs:
+        Tensor, with the same type and shape as the `input`.
+
+    Raises:
+        TypeError: If dtype of `input` is neither float16 nor float32.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> import numpy as np
+        >>> input = Tensor(np.array([1, 2, 3, 2, 1]), mindspore.float16)
+        >>> tanh = mint.nn.Tanh()
+        >>> output = tanh(input)
+        >>> print(output)
+        [0.7617 0.964  0.995  0.964  0.7617]
+    """
+
+    def __init__(self):
+        """Initialize Tanh."""
+        super(Tanh, self).__init__()
+
+    def construct(self, input):
+        return mint.nn.functional.tanh(input)
+
+
 __all__ = [
     'LogSigmoid',
     'SiLU',
+    'ELU',
+    'GLU',
+    'Tanh',
 ]

mindspore/mint/nn/layer/basic.py

@@ -0,0 +1,123 @@
+# 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.
+# ============================================================================
+"""activation layer for mint"""
+from __future__ import absolute_import
+from __future__ import division
+
+from mindspore import mint
+from mindspore.nn.cell import Cell
+from mindspore import _checkparam as validator
+
+
+class Dropout2d(Cell):
+    r"""
+    During training, randomly zeroes some channels of the input tensor with probability `p`
+    from a Bernoulli distribution (For a 4-dimensional tensor with a shape of :math:`NCHW`,
+    the channel feature map refers to a 2-dimensional feature map with the shape of :math:`HW`).
+
+    For example, the :math:`j\_th` channel of the :math:`i\_th` sample in the batched input is a to-be-processed
+    `2D` tensor input[i,j].
+    Each channel will be zeroed out independently on every forward call with probability `p` using samples
+    from a Bernoulli distribution.
+
+    `Dropout2d` can improve the independence between channel feature maps.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Refer to :func:`mindspore.mint.nn.functional.dropout2d` for more details.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> import numpy as np
+        >>> dropout = mint.nn.Dropout2d(p=0.5)
+        >>> x = Tensor(np.ones([2, 1, 2, 3]), mindspore.float32)
+        >>> output = dropout(x)
+        >>> print(output.shape)
+        (2, 1, 2, 3)
+    """
+
+    def __init__(self, p=0.5):
+        """Initialize Dropout2d."""
+        super(Dropout2d, self).__init__()
+        validator.check_float_range(p, 0.0, 1.0, validator.INC_BOTH, "p", self.cls_name)
+        self.p = p
+
+    def construct(self, x):
+        if not self.training or self.p == 0:
+            return x
+
+        return mint.nn.functional.dropout2d(x, self.p)
+
+
+class Flatten(Cell):
+    r"""
+    Flatten the input Tensor along dimensions from `start_dim` to `end_dim`.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        start_dim (int, optional): The first dimension to flatten. Default: ``1`` .
+        end_dim (int, optional): The last dimension to flatten. Default: ``-1`` .
+
+    Inputs:
+        - **input** (Tensor) - The input Tensor to be flattened.
+
+    Outputs:
+        Tensor. If no dimensions are flattened, returns the original `input`, otherwise return the flattened Tensor.
+
+    Raises:
+        TypeError: If `input` is not a Tensor.
+        TypeError: If `start_dim` or `end_dim` is not int.
+        ValueError: If `start_dim` is greater than `end_dim` after canonicalized.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> import numpy as np
+        >>> input = Tensor(np.array([[[1.2, 1.2], [2.1, 2.1]], [[2.2, 2.2], [3.2, 3.2]]]), mindspore.float32)
+        >>> net = mint.nn.Flatten()
+        >>> output = net(input)
+        >>> print(output)
+        [[1.2 1.2 2.1 2.1]
+         [2.2 2.2 3.2 3.2]]
+        >>> print(f"before flatten the x shape is {input.shape}")
+        before flatten the input shape is  (2, 2, 2)
+        >>> print(f"after flatten the output shape is {output.shape}")
+        after flatten the output shape is (2, 4)
+    """
+
+    def __init__(self, start_dim=1, end_dim=-1):
+        """Initialize Flatten."""
+        super(Flatten, self).__init__()
+        self.start_dim = start_dim
+        self.end_dim = end_dim
+
+    def construct(self, input):
+        return mint.nn.functional.flatten(input, self.start_dim, self.end_dim)
+
+
+__all__ = [
+    'Dropout2d',
+    'Flatten',
+]