mindspore 2.0.0rc1__cp38-none-any.whl → 2.2.0__cp38-none-any.whl

mindspore/nn/wrap/loss_scale.py

@@ -15,9 +15,11 @@
 """Loss scale cell for loss scale training."""
 from __future__ import absolute_import
 
+import os
 import mindspore.context as context
 from mindspore.context import ParallelMode
 from mindspore.parallel._utils import _get_enable_parallel_optimizer
+from mindspore import nn
 from mindspore.nn.wrap.cell_wrapper import TrainOneStepCell
 from mindspore.nn.cell import Cell
 from mindspore.common import Tensor
@@ -29,6 +31,7 @@ from mindspore.ops import composite as C
 from mindspore.ops import operations as P
 from mindspore.common import dtype as mstype
 from mindspore.common.api import jit
+from mindspore._c_expression import MSContext
 
 _grad_scale = C.MultitypeFuncGraph("grad_scale")
 reciprocal = P.Reciprocal()
@@ -59,6 +62,28 @@ def _tensor_grad_overflow_row_tensor(grad):
     return grad_overflow(grad.values)
 
 
+_ascend_grad_overflow = C.MultitypeFuncGraph("_ascend_grad_overflow")
+ascend_grad_overflow = P.IsFinite()
+
+
+@_ascend_grad_overflow.register("Tensor")
+def _tensor_ascend_grad_overflow(grad):
+    status = ascend_grad_overflow(grad)
+    base = Tensor(1.0, dtype=mstype.float32)
+    output = base - status.all()
+    output = P.Reshape()(output, ((1,)))
+    return output
+
+
+@_ascend_grad_overflow.register("RowTensor")
+def _tensor_ascend_grad_overflow_row_tensor(grad):
+    status = ascend_grad_overflow(grad.values)
+    base = Tensor(1.0, dtype=mstype.float32)
+    output = base - status.all()
+    output = P.Reshape()(output, ((1,)))
+    return output
+
+
 class DynamicLossScaleUpdateCell(Cell):
     r"""
     Dynamic Loss scale update cell.
@@ -88,8 +113,8 @@ class DynamicLossScaleUpdateCell(Cell):
 
     Examples:
         >>> import numpy as np
-        >>> from mindspore import Tensor, Parameter, nn
-        >>> import mindspore.ops as ops
+        >>> import mindspore
+        >>> from mindspore import Tensor, Parameter, nn, ops
         >>>
         >>> class Net(nn.Cell):
         ...     def __init__(self, in_features, out_features):
@@ -142,6 +167,13 @@ class DynamicLossScaleUpdateCell(Cell):
 
         Returns:
             float, the loss scale value.
+
+        Examples:
+            >>> from mindspore import nn
+            >>> manager = nn.DynamicLossScaleUpdateCell(loss_scale_value=212, scale_factor=2, scale_window=1000)
+            >>> output = manager.get_loss_scale()
+            >>> print(output)
+            212
         """
         return self.loss_scale_value
 
@@ -186,6 +218,7 @@ class FixedLossScaleUpdateCell(Cell):
 
     Examples:
         >>> import numpy as np
+        >>> import mindspore
         >>> from mindspore import Tensor, Parameter, nn, ops
         >>>
         >>> class Net(nn.Cell):
@@ -221,6 +254,13 @@ class FixedLossScaleUpdateCell(Cell):
 
         Returns:
             float, the loss scale value.
+
+        Examples:
+            >>> from mindspore import nn
+            >>> manager = nn.FixedLossScaleUpdateCell(loss_scale_value=212)
+            >>> output = manager.get_loss_scale()
+            >>> print(output)
+            212
         """
         return self.loss_scale_value
 
@@ -265,7 +305,6 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         >>> import numpy as np
         >>> import mindspore
         >>> from mindspore import Tensor, Parameter, nn, ops
-        >>> from mindspore import dtype as mstype
         >>>
         >>> class Net(nn.Cell):
         ...     def __init__(self, in_features, out_features):
@@ -281,14 +320,21 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         >>> size, in_features, out_features = 16, 16, 10
         >>> #1) when the type of scale_sense is Cell:
         >>> net = Net(in_features, out_features)
-        >>> loss = nn.MSELoss()
+        >>> loss_fn = nn.MSELoss()
         >>> optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
-        >>> net_with_loss = nn.WithLossCell(net, loss)
-        >>> manager = nn.DynamicLossScaleUpdateCell(loss_scale_value=2**12, scale_factor=2, scale_window=1000)
-        >>> train_network = nn.TrainOneStepWithLossScaleCell(net_with_loss, optimizer, scale_sense=manager)
+        >>> net_with_loss = nn.WithLossCell(net, loss_fn)
         >>> input = Tensor(np.ones([out_features, in_features]), mindspore.float32)
         >>> labels = Tensor(np.ones([out_features,]), mindspore.float32)
-        >>> output = train_network(input, labels)
+        >>> loss = net_with_loss(input, labels)
+        >>> manager = nn.DynamicLossScaleUpdateCell(loss_scale_value=2**12, scale_factor=2, scale_window=1000)
+        >>> train_network = nn.TrainOneStepWithLossScaleCell(net_with_loss, optimizer, scale_sense=manager)
+        >>> status = Tensor([0] * 8, mindspore.int32)
+        >>> scaling_sens = train_network.scale_sense
+        >>> scaling_sens_filled = ops.ones_like(loss) * ops.cast(scaling_sens, ops.dtype(loss))
+        >>> grads = train_network.grad(train_network.network, train_network.weights)(input, labels, scaling_sens_filled)
+        >>> grads = train_network.grad_reducer(grads)
+        >>> cond = train_network.get_overflow_status(status, grads)
+        >>> overflow = train_network.process_loss_scale(cond)
         >>>
         >>> #2) when the type of scale_sense is Tensor:
         >>> net = Net(in_features, out_features)
@@ -297,12 +343,14 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         >>> net_with_loss = nn.WithLossCell(net, loss)
         >>> inputs = Tensor(np.ones([size, in_features]).astype(np.float32))
         >>> label = Tensor(np.zeros([size, out_features]).astype(np.float32))
-        >>> scaling_sens = Tensor([1024], dtype=mstype.float32)
+        >>> scaling_sens = Tensor([1024], dtype=mindspore.float32)
         >>> train_network = nn.TrainOneStepWithLossScaleCell(net_with_loss, optimizer, scale_sense=scaling_sens)
+        >>> scaling_sens = Tensor([1], dtype=mstype.float32)
+        >>> train_network.set_sense_scale(scaling_sens)
         >>> output = train_network(inputs, label)
         >>>
         >>> # update scaling sens and train the network
-        >>> scaling_sens = Tensor([1], dtype=mstype.float32)
+        >>> scaling_sens = Tensor([1], dtype=mindspore.float32)
         >>> train_network.set_sense_scale(scaling_sens)
         >>> output = train_network(inputs, label)
     """
@@ -312,13 +360,19 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         self.base = Tensor(1, mstype.float32)
         self.base0 = Tensor(0, mstype.int32)
         self.reduce_sum = P.ReduceSum(keep_dims=False)
-        self.reduce_any = P.ReduceAny(keep_dims=False)
+        self.reduce_all = P.ReduceAll(keep_dims=False)
         self.less_equal = P.LessEqual()
-        self.not_equal = P.NotEqual()
+        self.equal = P.Equal()
+        self.logic_not = P.LogicalNot()
         self.allreduce = P.AllReduce()
         self.is_distributed = (self.parallel_mode != ParallelMode.STAND_ALONE)
         self.gpu_target = (context.get_context("device_target") == "GPU")
+        self.ascend_910a_target = (MSContext.get_instance().get_ascend_soc_version() == 'ascend910')
+        self.ascend_910b_target = (MSContext.get_instance().get_ascend_soc_version() == 'ascend910b')
         self.loss_scaling_manager = None
+        self._ascend910b_check_overflow_status_mode = os.environ.get('MS_ASCEND_CHECK_OVERFLOW_MODE')
+
+
         if isinstance(scale_sense, Cell):
             self.loss_scaling_manager = scale_sense
             self.scale_sense = Parameter(Tensor(scale_sense.get_loss_scale(), dtype=mstype.float32),
@@ -335,13 +389,13 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
                             "the 'scale_sense' must be Cell or Tensor, but got 'scale_sense' type: {}."
                             .format(type(scale_sense)))
         self.enable_tuple_broaden = True
+        self._get_attr_from_cell(network)
 
     def construct(self, *inputs):
         weights = self.weights
         loss = self.network(*inputs)
         scaling_sens = self.scale_sense
-
-        status, scaling_sens = self.start_overflow_check(loss, scaling_sens)
+        status = Tensor([0] * 8, mstype.int32)
 
         scaling_sens_filled = C.ones_like(loss) * F.cast(scaling_sens, F.dtype(loss))
         grads = self.grad(self.network, weights)(*inputs, scaling_sens_filled)
@@ -396,13 +450,68 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
             is cleaned up when the function returns.
         """
         status = Tensor([0] * 8, mstype.int32)
-        if not self.gpu_target:
+        if self.ascend_910a_target or (self.ascend_910b_target and \
+                                       self._ascend910b_check_overflow_status_mode != "INFNAN_MODE"):
             status = F.depend(status, pre_cond)
             # clear overflow buffer
             clear_status = NPUClearFloatStatusV2()(status)
             compute_input = F.depend(compute_input, clear_status)
         return status, compute_input
 
+    def _check_overflow_status_on_infnan_mode(self, grad_overflow_check_func, compute_output):
+        """check overflow status on infnan mode."""
+        flag_sum = self.hyper_map(F.partial(grad_overflow_check_func), compute_output)
+        flag_sum = P.AddN()(flag_sum)
+        # convert flag_sum to scalar
+        flag_sum = P.Reshape()(flag_sum, (()))
+        return flag_sum
+
+    def _get_distributed_overflow_status_on_infnan_mode(self, grad_overflow_check_func, compute_output):
+        """converge the distributed overflow status on infnan mode."""
+        flag_sum = self._check_overflow_status_on_infnan_mode(grad_overflow_check_func, compute_output)
+
+        if self.is_distributed:
+            # sum overflow flag over devices
+            flag_reduce = self.allreduce(flag_sum)
+            overflow = self.less_equal(self.base, flag_reduce)
+        else:
+            overflow = self.less_equal(self.base, flag_sum)
+        return overflow
+
+    def _get_gpu_overflow_status(self, compute_output):
+        """get overflow status of gpu."""
+        overflow = self._get_distributed_overflow_status_on_infnan_mode(_grad_overflow, compute_output)
+        return overflow
+
+    def _get_ascend_overflow_status_on_infnan_mode(self, compute_output):
+        """get overflow status of ascend on infnan mode."""
+        overflow = self._get_distributed_overflow_status_on_infnan_mode(_ascend_grad_overflow, compute_output)
+        return overflow
+
+    def _get_ascend_overflow_status_on_saturation_mode(self, status, compute_output):
+        """get overflow status of ascend on saturation mode"""
+        status = F.depend(status, compute_output)
+        get_status = NPUGetFloatStatusV2()(status)
+
+        if self.is_distributed:
+            # sum overflow flag over devices
+            flag_reduce = self.allreduce(get_status)
+            # get_status not equal to [0]*8 means overflow
+            flag = self.equal(self.base0, flag_reduce)
+            status = F.depend(status, flag)
+            # distributed needs to skip allreduce to avoid its overflow affecting the next step
+            clear_status = NPUClearFloatStatusV2()(status)
+            flag = F.depend(flag, clear_status)
+            overall_finite = self.reduce_all(flag)
+        else:
+            status = F.depend(status, get_status)
+            clear_status = NPUClearFloatStatusV2()(status)
+            get_status = F.depend(get_status, clear_status)
+            flag = self.equal(self.base0, get_status)
+            overall_finite = self.reduce_all(flag)
+        overflow = self.logic_not(overall_finite)
+        return overflow
+
     @jit
     def get_overflow_status(self, status, compute_output):
         """
@@ -420,37 +529,15 @@ class TrainOneStepWithLossScaleCell(TrainOneStepCell):
         Returns:
             bool, whether the overflow occurs or not.
         """
-        if not self.gpu_target:
-            status = F.depend(status, compute_output)
-            get_status = NPUGetFloatStatusV2()(status)
-
-            if self.is_distributed:
-                # sum overflow flag over devices
-                flag_reduce = self.allreduce(get_status)
-                # get_status not equal to [0]*8 means overflow
-                flag = self.not_equal(self.base0, flag_reduce)
-                status = F.depend(status, flag)
-                # distributed needs to skip allreduce to avoid its overflow affecting the next step
-                clear_status = NPUClearFloatStatusV2()(status)
-                flag = F.depend(flag, clear_status)
-            else:
-                status = F.depend(status, get_status)
-                clear_status = NPUClearFloatStatusV2()(status)
-                get_status = F.depend(get_status, clear_status)
-                flag = self.not_equal(self.base0, get_status)
-            overflow = self.reduce_any(flag)
-        else:
-            flag_sum = self.hyper_map(F.partial(_grad_overflow), compute_output)
-            flag_sum = P.AddN()(flag_sum)
-            # convert flag_sum to scalar
-            flag_sum = P.Reshape()(flag_sum, (()))
-
-            if self.is_distributed:
-                # sum overflow flag over devices
-                flag_reduce = self.allreduce(flag_sum)
-                overflow = self.less_equal(self.base, flag_reduce)
+        if self.gpu_target:
+            overflow = self._get_gpu_overflow_status(compute_output)
+        elif self.ascend_910b_target:
+            if self._ascend910b_check_overflow_status_mode != "INFNAN_MODE":
+                overflow = self._get_ascend_overflow_status_on_saturation_mode(status, compute_output)
             else:
-                overflow = self.less_equal(self.base, flag_sum)
+                overflow = self._get_ascend_overflow_status_on_infnan_mode(compute_output)
+        else:  # ascend_910a_target
+            overflow = self._get_ascend_overflow_status_on_saturation_mode(status, compute_output)
         return overflow
 
     def process_loss_scale(self, overflow):
@@ -493,7 +580,7 @@ def tensor_shard_grad_scale_pipeline(scale, grad, accu_grad):
     return new_grad
 
 
-class _TrainPipelineWithLossScaleCell(TrainOneStepCell):
+class _TrainGradAccuWithLossScaleCell(TrainOneStepCell):
     """
     Append an optimizer to the training network after that the construct
     function can be called to create the backward graph.
@@ -504,14 +591,14 @@ class _TrainPipelineWithLossScaleCell(TrainOneStepCell):
         scale_sense (Cell): Cell to do the loss scale.
     """
     def __init__(self, network, optimizer, scale_sense):
-        super(_TrainPipelineWithLossScaleCell, self).__init__(network, optimizer, sens=None)
+        super(_TrainGradAccuWithLossScaleCell, self).__init__(network, optimizer, sens=None)
         self.network = network
         self.network.add_flags(defer_inline=True)
         self.weights = optimizer.parameters
         self.accu_grads = self.weights.clone(prefix="accu_grads", init="zeros")
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)
-        self.grad_reducer = F.identity
+        self.grad_reducer = nn.Identity()
         self.degree = 1
         self.cast = P.Cast()
         self.alloc_status = P.NPUAllocFloatStatus()
@@ -547,12 +634,12 @@ class _TrainPipelineWithLossScaleCell(TrainOneStepCell):
         scaling_sens = self.scale_sense
         init = self.alloc_status()
         scaling_sens_filled = C.ones_like(loss) * F.cast(scaling_sens, F.dtype(loss))
+        scaling_sens_filled = F.depend(scaling_sens_filled, self.clear_before_grad(init))
         grads = self.grad(self.network, self.weights)(*inputs, scaling_sens_filled)
         init = F.depend(init, grads)
         get_status = self.get_status(init)
         init = F.depend(init, get_status)
         flag_sum = self.reduce_sum(init, (0,))
-        loss = F.depend(loss, self.clear_before_grad(init))
         if self.opt_shard:
             grads = self.grad_reducer(grads)
             grads = self.hyper_map(F.partial(shard_grad_scale, scaling_sens * self.degree), grads, self.accu_grads)

mindspore/numpy/__init__.py

@@ -33,7 +33,7 @@ from mindspore.numpy.array_ops import (transpose, expand_dims, squeeze, rollaxis
                                        tile, broadcast_to, broadcast_arrays, roll, append, split, vsplit,
                                        flip, flipud, fliplr, hsplit, dsplit, take_along_axis, take, repeat,
                                        rot90, select, array_split, choose, size, array_str, apply_along_axis,
-                                       piecewise, unravel_index, apply_over_axes)
+                                       piecewise, unravel_index, apply_over_axes, argwhere, intersect1d)
 from mindspore.numpy.array_creations import copy_ as copy
 from mindspore.numpy.array_creations import (array, asarray, asfarray, ones, zeros, full, randn, rand, randint,
                                              linspace, logspace, eye, identity, empty, empty_like, arange,
@@ -62,7 +62,7 @@ from mindspore.numpy.math_ops import (mean, inner, add, subtract, multiply, divi
 from mindspore.numpy.logic_ops import (not_equal, less_equal, less, greater_equal, greater, equal, isfinite,
                                        isnan, isinf, isposinf, isneginf, isscalar, logical_and, logical_not,
                                        logical_or, logical_xor, in1d, isin, isclose, signbit, sometrue,
-                                       array_equal, array_equiv)
+                                       array_equal, array_equiv, setdiff1d)
 
 mod = remainder
 fabs = absolute
@@ -82,7 +82,7 @@ array_ops_module = ['transpose', 'expand_dims', 'squeeze', 'rollaxis', 'swapaxes
                     'tile', 'broadcast_to', 'broadcast_arrays', 'append', 'roll', 'split', 'vsplit',
                     'flip', 'flipud', 'fliplr', 'hsplit', 'dsplit', 'take_along_axis', 'take',
                     'repeat', 'rot90', 'select', 'array_split', 'choose', 'size', 'array_str',
-                    'apply_along_axis', 'piecewise', 'unravel_index', 'apply_over_axes']
+                    'apply_along_axis', 'piecewise', 'unravel_index', 'apply_over_axes', 'argwhere', 'intersect1d']
 
 array_creations_module = ['array', 'asarray', 'asfarray', 'ones', 'zeros', 'full', 'randn', 'rand',
                           'randint', 'arange', 'linspace', 'logspace', 'eye', 'identity', 'empty',
@@ -112,7 +112,7 @@ math_module = ['mean', 'inner', 'add', 'subtract', 'multiply', 'divide', 'true_d
 logic_module = ['not_equal', 'less_equal', 'less', 'greater_equal', 'greater', 'equal', 'isfinite',
                 'isnan', 'isinf', 'isposinf', 'isneginf', 'isscalar', 'logical_and', 'logical_not',
                 'logical_or', 'logical_xor', 'in1d', 'isin', 'isclose', 'signbit', 'sometrue',
-                'array_equal', 'array_equiv']
+                'array_equal', 'array_equiv', 'setdiff1d']
 
 __all__ = array_ops_module + array_creations_module + math_module + logic_module + numeric_types