mindspore 2.3.0rc1__cp38-cp38-manylinux1_x86_64.whl → 2.3.0rc2__cp38-cp38-manylinux1_x86_64.whl

mindspore/common/recompute.py

@@ -0,0 +1,247 @@
+# Copyright 2023 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.
+# ============================================================================
+"""Defines other operators with functional form."""
+
+from collections import OrderedDict
+from types import MethodType
+from mindspore import log as logger
+from mindspore.nn.cell import Cell
+from mindspore import context
+from mindspore.common.tensor import Tensor
+from mindspore.ops.composite import GradOperation
+from mindspore.common._register_for_recompute import recompute_registry
+from mindspore.common.api import _pynative_executor
+from mindspore.nn.generator import get_rng_state, set_rng_state
+
+
+class _WrapCell(Cell):
+    """
+    The warp cell is used by recompute cell,
+    which can set mixed precision to warp cell
+    """
+
+    def __init__(self, function):
+        super(_WrapCell, self).__init__()
+        self.function = function
+
+    def construct(self, *args, **kwargs):
+        return self.function(*args, **kwargs)
+
+
+class _RecomputeCell(Cell):
+    """
+    Recompute cell, given the sub block, this cell will recompute the block, rather than
+    storing the intermediate activation computed in forward pass, we will recompute it in backward pass.
+    Note:
+     - RecomputeCell now only support pynative mode.
+     - When use recompute function, block object should not decorated by @jit.
+    """
+
+    def __init__(self, block):
+        """Initialize Recompute cell."""
+        super(_RecomputeCell, self).__init__()
+        self.args = []
+        self.kwargs = []
+        self.wrap_cell = _WrapCell(block)
+        self.net = block
+        self.internal_params = []
+        self.save_rng_state = False
+        self.cpu_rng_state = None
+        self._add_attr("is_cell_recompute", "True")
+        self.grad = GradOperation(get_all=True, get_by_list=True, sens_param=True)
+        self.init_mixed_precision_type(block)
+
+    def construct(self, *args, **kwargs):
+        _check_input_args_validate(self.net, args)
+        self.args.append(args)
+        self.kwargs.append(kwargs)
+        self.save_rng_state = kwargs.pop("save_rng_state", True)
+        if self.save_rng_state:
+            self.cpu_rng_state = get_rng_state()
+        return self.net(*args, **kwargs)
+
+    def bprop(self, *args):
+        """
+        Custom grad method for recompute
+        :param args:
+        :return: input grad and weight grads
+        """
+        grad_input = args[-1]
+        input_args = self.args[-1]
+        kwargs = self.kwargs[-1]
+        self.args.pop()
+        self.kwargs.pop()
+        if kwargs:
+            input_args = list(input_args) + list(kwargs.values())
+        try:
+            pre_rng_state = get_rng_state()
+            set_rng_state(*self.cpu_rng_state)
+            _pynative_executor.set_is_run_recompute(True)
+            grads = self.grad(self.net, self.internal_params)(*input_args, grad_input)
+            _pynative_executor.set_is_run_recompute(False)
+            set_rng_state(*pre_rng_state)
+        except Exception as err:
+            _pynative_executor.clear_res()
+            raise err
+        weights = OrderedDict()
+        input_grads = list(grads[0])
+        _padding_input_grads(input_args, input_grads)
+        for i, param in enumerate(self.internal_params):
+            weights[param] = grads[1][i]
+        return tuple(input_grads), weights
+
+    def init_mixed_precision_type(self, block):
+        """
+        init mix precision
+        :param block:
+        :return:
+        """
+        if isinstance(block, Cell):
+            # To avoid sub cell same name
+            block.check_names_and_refresh_name()
+            self.internal_params = block.trainable_params()
+            return
+        if isinstance(block, MethodType) and isinstance(block.__self__, Cell):
+            # To avoid sub cell same name
+            block.__self__.check_names_and_refresh_name()
+            self.internal_params = block.__self__.trainable_params()
+            self.wrap_cell.set_mixed_precision_type(block.__self__.get_mixed_precision_type())
+            self.net = self.wrap_cell
+        else:
+            raise TypeError("For Recompute cell, it not support FunctionType function, "
+                            "only support Cell object or MethodType function!")
+
+
+def _check_input_args_validate(block, args):
+    """
+    Check recompute input args validate
+    :param args:
+    :return:
+    """
+    if not any([isinstance(arg, Tensor) for arg in args]):
+        logger.warning("None of the inputs of function are tensors, which not need use recompute!")
+    for arg in args:
+        if isinstance(arg, (tuple, list)):
+            for data in arg:
+                if isinstance(data, Tensor):
+                    logger.info("For recompute block {}, tensor input in Tuple or list "
+                                "will not calculate grads!".format(block))
+                    break
+
+
+def _padding_input_grads(args, input_grads):
+    """
+    Padding input grads to same as input args
+    :param args:
+    :param input_grads:
+    :return:
+    """
+    for i, arg in enumerate(args):
+        if isinstance(arg, (list, tuple)):
+            if all([not isinstance(data, Tensor) for data in arg]):
+                input_grads.insert(i, None)
+            else:
+                # None is placeholder
+                grads = [None for data in arg]
+                input_grads.insert(i, grads)
+        elif not isinstance(arg, Tensor):
+            input_grads.insert(i, None)
+    if len(args) != len(input_grads):
+        raise ValueError("For recompute cell, the input grads size should be same as input args size: {}, "
+                         "but got {}".format(len(args), len(input_grads)))
+
+
+def _check_validation(block):
+    if not isinstance(block, Cell):
+        raise TypeError("Recompute function now only support block which inherited from Cell!")
+    if context.get_context("mode") != context.PYNATIVE_MODE:
+        raise AssertionError("Recompute function now only support pynative mode, you can use "
+                             "Cell.recompute() in graph mode.")
+    if block.construct.__code__.co_name == "staging_specialize":
+        logger.warning('Block\'s construct method decorated by @jit that recompute '
+                       'function will not come into effect.')
+
+
+def recompute(block, *args, **kwargs):
+    r"""
+    This function is used to reduce memory, when run block, rather than
+    storing the intermediate activation computed in forward pass, we will recompute it in backward pass.
+
+    Note:
+        - Recompute function only support block which inherited from Cell object.
+        - This function interface now only support pynative mode. you can use Cell.recompute interface
+          in graph mode.
+        - When use recompute function, block object should not decorated by @jit.
+
+    Args:
+        block (Cell): Block to be recompute.
+        args(tuple): Inputs for block object to run forward pass.
+        kwargs(dict): Optional input for recompute function.
+
+    Returns:
+        Same as return type of block.
+
+    Raises:
+        TypeError: If `block` is not Cell object.
+        AssertionError: If execute mode is not PYNATIVE_MODE.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore.nn as nn
+        >>> import mindspore.ops as ops
+        >>> from mindspore import Tensor, recompute
+        >>> class MyCell(nn.Cell):
+        ...     def __init__(self):
+        ...         super(MyCell, self).__init__(auto_prefix=False)
+        ...         self.conv = nn.Conv2d(2, 2, 2, has_bias=False, weight_init='ones')
+        ...         self.relu = ops.ReLU()
+        ...
+        ...     def construct(self, x):
+        ...         y = recompute(self.conv, x)
+        ...         return self.relu(y)
+        >>> inputs = Tensor(np.ones([2, 2, 2, 2]).astype(np.float32) * 2)
+        >>> my_net = MyCell()
+        >>> grad = ops.grad(my_net)(inputs)
+        >>> print(grad)
+        [[[[2. 4.]
+           [4. 8.]]
+          [[2. 4.]
+           [4. 8.]]]
+         [[[2. 4.]
+           [4. 8.]]
+          [[2. 4.]
+           [4. 8.]]]]
+    """
+
+    _check_validation(block)
+    return _RecomputeCell(block)(*args, **kwargs)
+
+
+def recompute_generator(block):
+    """
+    generator of recompute object.
+    :param block:
+    :return:
+    """
+    return _RecomputeCell(block)
+
+
+recompute_registry.register(recompute_generator)
+
+__all__ = ['recompute']
+__all__.sort()

mindspore/common/sparse_tensor.py

@@ -226,7 +226,7 @@ class COOTensor(COOTensor_):
 
     Common arithmetic operations include: addition (+), subtraction (-), multiplication (*),
     and division (/). For details about operations supported by `COOTensor`, see
-    `operators <https://www.mindspore.cn/docs/en/r2.3.q1/note/static_graph_syntax_support.html#operators>`_.
+    `operators <https://www.mindspore.cn/docs/en/master/note/static_graph_syntax_support.html#operators>`_.
 
     .. warning::
         - This is an experimental API that is subject to change or deletion.
@@ -653,7 +653,7 @@ class CSRTensor(CSRTensor_):
 
     Common arithmetic operations include: addition (+), subtraction (-), multiplication (*),
     and division (/). For details about operations supported by `CSRTensor`, see
-    `operators <https://www.mindspore.cn/docs/en/r2.3.q1/note/static_graph_syntax_support.html#operators>`_.
+    `operators <https://www.mindspore.cn/docs/en/master/note/static_graph_syntax_support.html#operators>`_.
 
     .. warning::
         - This is an experimental API that is subjected to change.

mindspore/common/symbol.py

@@ -22,7 +22,7 @@ class Symbol:
     Symbol is a data structure to indicate the symbolic info of shape.
 
     For dynamic shape networks, compared with only setting the unknown dimensions ( ``None`` ) in `Tensor` , providing
-    more symbolic shape info can help the framework better optimize the computation graph, to improve the performce of
+    more symbolic shape info can help the framework better optimize the computation graph, to improve the performance of
     network execution.
 
     Args:

mindspore/common/tensor.py

@@ -27,12 +27,13 @@ from mindspore.common.seed import get_seed
 from mindspore import context
 from mindspore import log as logger
 from mindspore.common import dtype as mstype
+from mindspore.common.hook_handle import _TensorHookHandle
 
 from mindspore.common._utils import get_slice_num
 from mindspore.common._register_for_tensor import tensor_operator_registry
 from mindspore._c_expression import Tensor as Tensor_
 from mindspore import _checkparam as validator
-from mindspore._checkparam import check_is_number, is_stub_tensor
+from mindspore._checkparam import check_is_number, is_stub_tensor, check_hook_fn
 from mindspore._check_jit_forbidden_api import jit_forbidden_register
 from mindspore.common.symbol import Symbol
 
@@ -83,11 +84,11 @@ def tensor(input_data=None, dtype=None, shape=None, init=None, internal=False, c
     based on the `dtype` argument.
 
     Please refer to `Creating and Using Tensor
-    <https://www.mindspore.cn/docs/en/r2.3.q1/note/static_graph_syntax_support.html#mindspore-user-defined-data-types>`_ .
+    <https://www.mindspore.cn/docs/en/master/note/static_graph_syntax_support.html#mindspore-user-defined-data-types>`_ .
 
     The difference between it and the Tensor class is that it adds
     `Annotation
-    <https://www.mindspore.cn/docs/en/r2.3.q1/design/dynamic_graph_and_static_graph.html?#annotation-type>`_
+    <https://www.mindspore.cn/docs/en/master/design/dynamic_graph_and_static_graph.html?#annotation-type>`_
     which can prevent the generation of AnyType compared to the Tensor class.
 
     The arguments and return values are the same as the Tensor class. Also see: :class:`mindspore.Tensor`.
@@ -1226,6 +1227,54 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         """
         return Tensor_.storage_offset(self)
 
+    def register_hook(self, hook_fn):
+        """
+        Registers a backward hook for tensor.
+
+        Note:
+            - The `register_backward_hook(hook_fn)` does not work in graph mode or functions decorated with 'jit'.
+            - The 'hook_fn' must be defined as the following code. `grad` is the gradient passed to the tensor,
+              which may be modified by returning a new output gradient.
+            - The 'hook_fn' should have the following signature:
+              hook_fn(grad) -> New output gradient, but can not return None or not set return value.
+
+        Args:
+            hook_fn (function): Python function. Tensor backward hook function.
+
+        Returns:
+            A handle corresponding to the `hook_fn` . The handle can be used to remove the added `hook_fn` by calling
+            `handle.remove()` .
+
+        Raises:
+            TypeError: If the `hook_fn` is not a function of python.
+
+        Supported Platforms:
+            ``Ascend`` ``GPU`` ``CPU``
+
+        Examples:
+            >>> import mindspore as ms
+            >>> from mindspore import Tensor
+            >>> ms.set_context(mode=ms.PYNATIVE_MODE)
+            >>> def hook_fn(grad):
+            ...     return grad * 2
+            ...
+            >>> def hook_test(x, y):
+            ...     z = x * y
+            ...     z.register_hook(hook_fn)
+            ...     z = z * y
+            ...     return z
+            ...
+            >>> ms_grad = ms.grad(hook_test, grad_position=(0,1))
+            >>> output = ms_grad(Tensor(1, ms.float32), Tensor(2, ms.float32))
+            >>> print(output)
+            (Tensor(shape=[], dtype=Float32, value=8), Tensor(shape=[], dtype=Float32, value=6))
+        """
+        if not check_hook_fn("register_hook", hook_fn):
+            return _TensorHookHandle()
+        handle = _TensorHookHandle()
+        handle.id = Tensor_.register_hook(self, hook_fn)
+        return handle
+
     def flush_from_cache(self):
         """
         Flush cache data to host if tensor is cache enable.
@@ -1815,9 +1864,9 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
 
     def reverse(self, axis):
         """
-        For details, please refer to :func:`mindspore.ops.reverse`.
+        For details, please refer to :func:`mindspore.ops.flip`.
         """
-        return tensor_operator_registry.get('reverse')(self, axis)
+        return tensor_operator_registry.get('flip')(self, axis)
 
     def amax(self, axis=None, keepdims=False, *, initial=None, where=None):
         """
@@ -1839,11 +1888,11 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         """
         return tensor_operator_registry.get("reverse_sequence")(self, seq_lengths, seq_dim, batch_dim)
 
-    def prod(self, axis=None, keep_dims=False):
+    def prod(self, axis=None, keep_dims=False, dtype=None):
         """
         For details, please refer to :func:`mindspore.ops.prod`.
         """
-        return tensor_operator_registry.get('prod')(self, axis, keep_dims)
+        return tensor_operator_registry.get('prod')(self, axis, keep_dims, dtype)
 
     def select(self, condition, y):
         r"""
@@ -2720,7 +2769,7 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
             opt_shard_group(str): Optimizer shard group which is used in auto or semi auto parallel mode
                 to get one shard of a parameter's slice. For more information about optimizer parallel, please refer to:
                 `Optimizer Parallel
-                <https://www.mindspore.cn/tutorials/experts/en/r2.3.q1/parallel/optimizer_parallel.html>`_.
+                <https://www.mindspore.cn/tutorials/experts/en/master/parallel/optimizer_parallel.html>`_.
                 Default: ``None``.
 
         Returns:
@@ -3306,14 +3355,9 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
             >>> print(input_x.sum(axis=1))
             [10. 35.]
         """
-        if initial is not None and not isinstance(initial, (int, float, bool)):
-            raise TypeError(f"For Tensor.sum, initial must be int, float or bool, but got {type(initial)}.")
-        res = tensor_operator_registry.get("sum")(self, axis, keepdims)
-        if initial is not None:
-            res += initial
-        if dtype is not None:
-            res = res.astype(dtype)
-        return res
+        if initial is None:
+            return tensor_operator_registry.get("sum")(self, axis, keepdims, dtype=dtype)
+        return tensor_operator_registry.get("sum")(self, axis, keepdims, dtype=dtype) + initial
 
     def sum_to_size(self, *size):
         r"""
@@ -4393,14 +4437,15 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
         """
         return tensor_operator_registry.get('not_equal')(self, other)
 
-    def new_zeros(self, size, *, dtype=None):
+    def new_zeros(self, size, dtype=None):
         r"""
         Return a tensor of `size` filled with zeros.
 
-        Args:
-            size (Union[int, tuple, list]): An int, list or tuple of integers defining the output shape.
+        .. warning::
+            For argument `size`, Tensor type input will be deprecated in the future version.
 
-        Keyword Args:
+        Args:
+            size (Union[int, tuple, list, Tensor]): An int, list or tuple of integers defining the output shape.
             dtype (mindspore.dtype, optional): The desired dtype of the output tensor. If None, the returned tensor has
                 thesame dtype as `self`. Default: ``None``.
 
@@ -4408,7 +4453,7 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
             Tensor, the shape and dtype is defined above and filled with zeros.
 
         Raises:
-            TypeError: If `size` is not an int, list or tuple of integers.
+            TypeError: If `size` is neither an int nor an tuple/list/Tensor of int.
 
         Supported Platforms:
             ``Ascend`` ``GPU`` ``CPU``
@@ -4423,20 +4468,17 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
             [[0. 0.]
              [0. 0.]]
         """
-        validator.check_value_type('size', size, [list, int, tuple], 'Tensor.new_zeros')
-        if isinstance(size, list):
-            size = tuple(size)
-        _dtype = self.dtype if dtype is None else dtype
-        return tensor_operator_registry.get('zeros')(size, _dtype)
+        return tensor_operator_registry.get('zeros')(size, dtype)
 
-    def new_ones(self, size, *, dtype=None):
+    def new_ones(self, size, dtype=None):
         r"""
         Return a tensor of `size` filled with ones.
 
-        Args:
-            size (Union[int, tuple, list]): An int, list or tuple of integers defining the output shape.
+        .. warning::
+            For argument `size`, Tensor type input will be deprecated in the future version.
 
-        Keyword Args:
+        Args:
+            size (Union[int, tuple, list, Tensor]): An int, list or tuple of integers defining the output shape.
             dtype (mindspore.dtype, optional): The desired dtype of the output tensor. If None, the returned
                 tensor has the same dtype as `self`. Default: ``None``.
 
@@ -4444,7 +4486,7 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
             Tensor, the shape and dtype is defined above and filled with ones.
 
         Raises:
-            TypeError: If `size` is not an int, list or tuple of integers.
+            TypeError: If `size` is neither an int nor an tuple/list/Tensor of int.
 
         Supported Platforms:
             ``Ascend`` ``GPU`` ``CPU``
@@ -4459,11 +4501,7 @@ class Tensor(Tensor_, metaclass=_TensorMeta):
             [[1. 1.]
              [1. 1.]]
         """
-        validator.check_value_type('size', size, [list, int, tuple], 'Tensor.new_zeros')
-        if isinstance(size, list):
-            size = tuple(size)
-        _dtype = self.dtype if dtype is None else dtype
-        return tensor_operator_registry.get('ones')(size, _dtype)
+        return tensor_operator_registry.get('ones')(size, dtype)
 
     def sign(self):
         r"""

mindspore/communication/__init__.py

@@ -19,14 +19,14 @@ Note that the APIs in the following list need to preset communication environmen
 
 For the Ascend devices, users need to prepare the rank table, set rank_id and device_id.
 Please see the `rank table Startup
-<https://www.mindspore.cn/tutorials/experts/en/r2.3.q1/parallel/rank_table.html>`_
+<https://www.mindspore.cn/tutorials/experts/en/master/parallel/rank_table.html>`_
 for more details.
 
 For the GPU devices, users need to prepare the host file and mpi, please see the `mpirun Startup
-<https://www.mindspore.cn/tutorials/experts/en/r2.3.q1/parallel/mpirun.html>`_ .
+<https://www.mindspore.cn/tutorials/experts/en/master/parallel/mpirun.html>`_ .
 
 For the CPU device, users need to write a dynamic cluster startup script, please see the `Dynamic Cluster Startup
-<https://www.mindspore.cn/tutorials/experts/en/r2.3.q1/parallel/dynamic_cluster.html>`_ .
+<https://www.mindspore.cn/tutorials/experts/en/master/parallel/dynamic_cluster.html>`_ .
 """
 
 from mindspore.communication.management import GlobalComm, init, release, get_rank, \