mindspore 2.4.10__cp310-cp310-win_amd64.whl → 2.6.0rc1__cp310-cp310-win_amd64.whl

mindspore/mint/nn/layer/__init__.py

@@ -27,13 +27,30 @@ from mindspore.mint.nn.layer.normalization import BatchNorm1d
 from mindspore.mint.nn.layer.normalization import BatchNorm2d
 from mindspore.mint.nn.layer.normalization import BatchNorm3d
 from mindspore.mint.nn.layer.normalization import LayerNorm
+from mindspore.mint.nn.layer.normalization import SyncBatchNorm
 from mindspore.mint.nn.layer.activation import LogSigmoid
 from mindspore.mint.nn.layer.activation import SiLU
+from mindspore.mint.nn.layer.activation import Threshold
+from mindspore.mint.nn.layer.basic import Dropout2d
+from mindspore.mint.nn.layer.pooling import AdaptiveMaxPool1d
 from mindspore.mint.nn.layer.pooling import AdaptiveAvgPool1d
 from mindspore.mint.nn.layer.pooling import AdaptiveAvgPool2d
+from mindspore.mint.nn.layer.pooling import AdaptiveAvgPool3d
 
 
-__all__ = []
-__all__.extend(normalization.__all__)
-__all__.extend(activation.__all__)
-__all__.extend(pooling.__all__)
+__all__ = [
+    'GroupNorm',
+    'BatchNorm1d',
+    'BatchNorm2d',
+    'BatchNorm3d',
+    'LayerNorm',
+    'LogSigmoid',
+    'SiLU',
+    'Dropout2d',
+    'AdaptiveMaxPool1d',
+    'AdaptiveAvgPool1d',
+    'AdaptiveAvgPool2d',
+    'AdaptiveAvgPool3d',
+    'SyncBatchNorm',
+    'Threshold',
+]

mindspore/mint/nn/layer/_functions.py

@@ -0,0 +1,334 @@
+import mindspore
+from mindspore import Tensor
+from mindspore import context
+import mindspore.communication
+import mindspore.communication.comm_func
+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 TensorPy 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.communication.comm_func import all_gather_into_tensor as all_gather_into_tensor_dy
+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, _ = mindspore.communication.comm_func.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_dy(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

@@ -77,6 +77,55 @@ class SiLU(Cell):
         return mint.nn.functional.silu(x)
 
 
+class Sigmoid(Cell):
+    r"""
+    Applies sigmoid activation function element-wise.
+
+    Sigmoid function is defined as:
+
+    .. math::
+
+        \text{sigmoid}(x_i) = \frac{1}{1 + \exp(-x_i)},
+
+    where :math:`x_i` is the element of `x`.
+
+    Sigmoid Activation Function Graph:
+
+    .. image:: ../images/Sigmoid.png
+        :align: center
+
+    Inputs:
+        - **input** (Tensor) - `input` is :math:`x` in the preceding formula. Tensor of any dimension,
+          the data type is float16, float32, float64, complex64 or complex128.
+
+    Outputs:
+        Tensor, with the same type and shape as the `input`.
+
+    Raises:
+        TypeError: If dtype of `input` is not float16, float32, float64, complex64 or complex128.
+        TypeError: If `input` is not a Tensor.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, nn
+        >>> import numpy as np
+        >>> input = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float16)
+        >>> sigmoid = mint.nn.Sigmoid()
+        >>> output = sigmoid(input)
+        >>> print(output)
+        [0.2688  0.11914 0.5     0.881   0.7305 ]
+    """
+    def __init__(self):
+        """Initialize LogSigmoid."""
+        super(Sigmoid, self).__init__()
+
+    def construct(self, input):
+        return mint.nn.functional.sigmoid(input)
+
+
 class LogSigmoid(Cell):
     r"""
     Applies logsigmoid activation element-wise. The input is a Tensor with any valid shape.
@@ -84,7 +133,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 +176,233 @@ 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, and
+    `alpha` represents the smoothness of the ELU.
+
+    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``.
+        inplace (bool, optional): Whether to use inplace mode, the data type is bool. Default: ``False``.
+
+    Inputs:
+        - **input** (Tensor) - The input of ELU is a Tensor of any dimension.
+
+    Outputs:
+        Tensor, with the same shape and type as the `input`.
+
+    Raises:
+        RuntimeError: If the dtype of `input` is not float16, float32 or bfloat16.
+        TypeError: If the dtype of `alpha` is not 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, inplace=False):
+        """Initialize ELU."""
+        super(ELU, self).__init__()
+        self.alpha = alpha
+        self.inplace = inplace
+
+    def construct(self, input):
+        return mint.nn.functional.elu(input, self.alpha, self.inplace)
+
+
+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>`_ .
+
+    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:
+        >>> from mindspore import mint, Tensor
+        >>> glu = mint.nn.GLU()
+        >>> input = Tensor([[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]])
+        >>> output = glu(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)
+
+
+class Threshold(Cell):
+    r"""
+    Compute the Threshold activation function element-wise.
+
+    The Threshold is defined as:
+
+    .. math::
+        y =
+        \begin{cases}
+        x, &\text{ if } x > \text{threshold} \\
+        \text{value}, &\text{ otherwise }
+        \end{cases}
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Args:
+        threshold (Union[int, float]): The value of the threshold.
+        value (Union[int, float]): The value to replace with when element is less than threshold.
+        inplace (bool, optional): Whether to apply erasing inplace. Default: ``False``.
+
+    Inputs:
+        - **input** (Tensor) - The input Tensor.
+
+    Outputs:
+        Tensor, the same shape and data type as the input.
+
+    Raises:
+        TypeError: If `input` is not a Tensor.
+        TypeError: If `threshold` is not a float or an int.
+        TypeError: If `value` is not a float or an int.
+
+    Supported Platforms:
+        ``Ascend``
+
+    Examples:
+        >>> import mindspore
+        >>> from mindspore import Tensor, mint
+        >>> inputs = mindspore.Tensor([0.0, 2, 3], mindspore.float32)
+        >>> net = mint.nn.Threshold(1, 100)
+        >>> outputs = net(inputs)
+        >>> print(outputs)
+        [100.   2.   3.]
+    """
+
+    def __init__(self, threshold, value, inplace=False):
+        """Initialize Tanh."""
+        super(Threshold, self).__init__()
+        self.threshold = threshold
+        self.value = value
+        self.inplace = inplace
+
+    def construct(self, input):
+        return mint.nn.functional.threshold(input, self.threshold, self.value,
+                                            self.inplace)
+
 __all__ = [
     'LogSigmoid',
     'SiLU',
+    'ELU',
+    'GLU',
+    'Tanh',
+    'Threshold',
 ]