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

mindspore/train/amp.py

@@ -14,6 +14,8 @@
 # ============================================================================
 """Auto mixed precision."""
 from __future__ import absolute_import
+import inspect
+import types
 
 import mindspore as ms
 from mindspore import nn
@@ -61,8 +63,20 @@ AMP_BLACK_LIST = [
     nn.LayerNorm
 ]
 
+# Primitives in inner amp black list will not be converted in O2/O3
+_INNER_AMP_BLACK_LIST = []
+
 MS_AMP_BY_REWRITE = False
-_amp_cast_op = P.Cast
+
+
+def amp_cast(value, dtype):
+    """This function is used to insert cast operators for tensors during auto mixed precision."""
+    if isinstance(value, ms.Tensor) and value.dtype in mstype.float_type:
+        return P.Cast()(value, dtype)
+    return value
+
+_amp_cast_op = amp_cast
+
 
 class _OutputTo16(nn.Cell):
     """Wrap cell for amp. Cast network output back to float16."""
@@ -88,278 +102,185 @@ class _OutputTo32(nn.Cell):
         return F.mixed_precision_cast(mstype.float32, out)
 
 
-
-def _allow_mix_precision(node, allowed_list, dtype) -> bool:
+def _operator_need_cast(node, force_cast: bool, white_list=None, black_list=None) -> bool:
     """
-    Check whether current node need do mix precision. Follow conditions need to be satisfied:
-        1) Type of node is one of (Primitive, nn.Cell)
-        2) Node is not Cast Op
-        3) to_float(mindspore.float16) is not set in Cell
+    Check whether current node is a operator that need to be casted. Follow conditions need to be satisfied:
+        1) Type of node is CallPrimitive and type of instance is Primitive
+        2) Type of instance is not P.Cast
+        3) force_cast is True, which means one of upper layer cells is under casting
+        4) white_list exist and type of node is in white_list
+        5) black_list exist and type of node is in not black_list
     """
-    node_inst = node.get_instance()
-    if node_inst in allowed_list:
-        return True
-    if node.get_targets() is None:
+    if node.get_node_type() != ms.rewrite.NodeType.CallPrimitive:
         return False
-    if not issubclass(node.get_instance_type(), (Primitive, nn.Cell)):
+    if not inspect.isclass(node.get_instance_type()):
         return False
-    if isinstance(node_inst, _amp_cast_op):
+    if not issubclass(node.get_instance_type(), Primitive):
         return False
-    if issubclass(node.get_instance_type(), nn.Cell):
-        # if cell is already in allowed_list, it means to_float() is set by amp.
-        # if cell is not in allowed_list, but has to_float(),
-        # it means to_float() is set by user.
-        to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-        if hasattr(node_inst, to_float_flag) and getattr(node_inst, to_float_flag):
-            return False
-    allowed_list.append(node.get_instance())
-    return True
+    if issubclass(node.get_instance_type(), P.Cast):
+        return False
+    if node.get_instance_type() in _INNER_AMP_BLACK_LIST:
+        return False
+    if force_cast:
+        return True
+    if white_list is not None and node.get_instance_type() in white_list:
+        return True
+    if black_list is not None and node.get_instance_type() not in black_list:
+        return True
+    return False
 
 
-def _insert_cast_operator_process(node, dtype):
-    """insert cast for operators in white_list."""
-    dtype_str = "mindspore.bfloat16" if dtype == mstype.bfloat16 else "mindspore.float16"
-    new_cast_node = None
-    stree = node.get_symbol_tree()
-    # insert cast fp16/bf16 before the primitive operators
-    if issubclass(node.get_instance_type(), Primitive):
-        for idx, arg in enumerate(node.get_args()):
-            position = stree.before(node)
-            new_node = _amp_cast_op()
-            cast_args = ms.rewrite.ScopedValue.create_name_values([arg.value, dtype_str], [arg.scope, ""])
-            arg_provider = node.get_handler().get_arg_providers()[idx]
-            if arg_provider and len(arg_provider[0].get_target_users(arg_provider[1])) > 1:
-                cast_targets = [stree.unique_name(str(arg))]
-            else:
-                cast_targets = ms.rewrite.ScopedValue.create_name_values([arg.value], [arg.scope])
-            new_cast_node = ms.rewrite.Node.create_call_cell(new_node,
-                                                             targets=cast_targets,
-                                                             args=cast_args,
-                                                             name='incast_{}{}'.format(node.get_name(), idx))
-            stree.insert(position, new_cast_node)
-            node.set_arg_by_node(idx, new_cast_node)
-    # insert cast fp16/bf16 before the Cell operators
-    elif issubclass(node.get_instance_type(), nn.Cell):
-        node.get_instance().to_float(dtype)
-    # ignore if subclass is not one of (Primitive, nn.Cell)
-    else:
-        return
-
-    # insert cast float32 after the operators
-    position = stree.after(node)
-    new_node = _amp_cast_op()
-    cast_args = ms.rewrite.ScopedValue.create_name_values([node.get_targets()[0].value,
-                                                           "mindspore.float32"])
-    new_cast_node = ms.rewrite.Node.create_call_cell(new_node,
-                                                     targets=[node.get_targets()[0]],
-                                                     args=cast_args,
-                                                     name='outcast_{}'.format(node.get_name()))
-    # insert node & unique names
-    stree.insert(position, new_cast_node)
-    # update argument names
-    for user in node.get_users():
-        if user.get_name() == new_cast_node.get_name():
-            continue
-        for idx, arg in enumerate(user.get_args()):
-            if arg == node.get_targets()[0]:
-                user.set_arg_by_node(idx, new_cast_node)
-
-
-def _insert_cast_operator_white_list(stree, white_list, dtype):
-    """insert cast for operators in white_list."""
-    allowed_list = []
-    # Ignore if net called ".to_float(dtype)"
-    net = stree.get_handler().get_origin_network()
-    to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-    if isinstance(net, nn.Cell) and hasattr(net, to_float_flag) and getattr(net, to_float_flag):
-        return
-    node_list = []
-    node_list.extend(list(stree.nodes()))
-    while node_list:
-        node = node_list.pop()
-        if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
-            if MS_AMP_BY_REWRITE:
-                _insert_cast_for_cell_container(node, dtype, allowed_list, white_list=white_list)
-            for n in node.get_handler().node_list:
-                if n.get_node_type() == ms.rewrite.NodeType.Tree:
-                    _insert_cast_operator_white_list(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)),
-                                                     white_list, dtype)
-        elif node.get_node_type() == ms.rewrite.NodeType.Tree:
-            substree = ms.rewrite.TreeNodeHelper.get_sub_tree(node)
-            _insert_cast_operator_white_list(substree, white_list, dtype)
-        elif node.get_node_type() in [ms.rewrite.NodeType.CallFunction, ms.rewrite.NodeType.ControlFlow]:
-            if isinstance(node.get_handler(), ms.rewrite.node.NodeManager):
-                nodes = [ms.rewrite.Node(n) for n in node.get_handler().nodes()]
-                node_list.extend(nodes)
-        elif node.get_instance_type() in white_list and _allow_mix_precision(node, allowed_list, dtype):
-            _insert_cast_operator_process(node, dtype)
+def _precision_set_by_user(cell_inst: nn.Cell) -> bool:
+    """Check whether cell precision is set by user."""
+    for flag in ["fp32", "fp16", "bf16"]:
+        if hasattr(cell_inst, flag) and getattr(cell_inst, flag):
+            return True
+    return False
 
 
-def _insert_cast_for_cell_container(cell_container, dtype, allowed_list, *, white_list=None, black_list=None):
+def _net_need_cast(node, force_cast: bool, white_list=None, black_list=None) -> bool:
     """
-    Insert cast for cell containers.
-    Only one of white_list and black_list can be set.
+    Check whether current node is type of tree whose network needs to be casted. Follow conditions need to
+    be satisfied:
+        1) Type of node is Tree and type of instance is Cell
+        2) Cell.to_float(xxx) is not set by user
+        3) force_cast is True, which means one of upper layer networks is under casting
+        4) white_list exist and type of node is in white_list
+        5) black_list exist and type of node is in not black_list
     """
+    if node.get_node_type() != ms.rewrite.NodeType.Tree:
+        return False
+    if not inspect.isclass(node.get_instance_type()):
+        return False
+    if not issubclass(node.get_instance_type(), nn.Cell):
+        return False
+    if node.get_instance_type() in _INNER_AMP_BLACK_LIST:
+        return False
+    if _precision_set_by_user(node.get_instance()):
+        return False
+    if force_cast:
+        return True
+    if white_list is not None and node.get_instance_type() in white_list:
+        return True
+    if black_list is not None and node.get_instance_type() not in black_list:
+        return True
+    return False
 
-    class CastNet(nn.Cell):
-        """Cast net"""
-        def __init__(self, dtype):
-            super().__init__()
-            self.cast = _amp_cast_op()
-            self.dtype = dtype
-
-        def construct(self, x):
-            return self.cast(x, self.dtype)
-
-    cast_flag = False
-    current_node = None
-    stree = cell_container.get_symbol_tree()
-    for node in cell_container.get_handler().nodes():
-        current_node = ms.rewrite.Node(node)
-        if (white_list is not None and current_node.get_instance_type() in white_list) or \
-           (black_list is not None and current_node.get_instance_type() not in black_list) and \
-           (_allow_mix_precision(current_node, allowed_list, dtype)):
-            cast_flag = True
-            current_node.get_instance().to_float(dtype)
-        elif cast_flag:
-            # cast next node back to float32
-            current_node.get_instance().to_float(mstype.float32)
-            cast_flag = False
-    if cast_flag and current_node:
-        # if last node in cell_container is casted to fp16/bf16, insert a cast node to cast value back to fp32
-        cast_node = ms.rewrite.Node.create_call_cell(cell=CastNet(mstype.float32),
-                                                     args=[current_node.get_targets()[0]],
-                                                     targets=[current_node.get_targets()[0]],
-                                                     name=f"outcast_{cell_container.get_name()}")
-        stree.insert(stree.after(current_node), cast_node)
+
+def _insert_cast_for_operator(node, dtype):
+    """insert cast pair for node."""
+    dtype_str = "bfloat16" if dtype == mstype.bfloat16 else "float16"
+    stree = node.get_symbol_tree()
+    # insert cast fp16/bf16 for inputs of node
+    for idx, arg in enumerate(node.get_args()):
+        if arg.type != ms.rewrite.ValueType.NamingValue:
+            continue
+        incast_args = ms.rewrite.ScopedValue.create_name_values([arg.value, dtype_str], [arg.scope, "mindspore"])
+        arg_providers = node.get_arg_providers()
+        if not arg_providers or idx not in arg_providers or \
+            len(arg_providers[idx][0].get_target_users(arg_providers[idx][1])) > 1:
+            # create new target names when argument is used by other node
+            incast_targets = [stree.unique_name(f"{arg.value}_var")]
+        else:
+            incast_targets = ms.rewrite.ScopedValue.create_name_values([arg.value], [arg.scope])
+        incast_node = ms.rewrite.Node.create_call_function(_amp_cast_op, targets=incast_targets, args=incast_args)
+        stree.insert(stree.before(node), incast_node)
+        node.set_arg_by_node(idx, incast_node)
+    # insert cast fp32 for outputs of node
+    for _, target in enumerate(node.get_targets()):
+        if target.type != ms.rewrite.ValueType.NamingValue:
+            continue
+        outcast_args = ms.rewrite.ScopedValue.create_name_values([target.value, "float32"],
+                                                                 [target.scope, "mindspore"])
+        outcast_targets = ms.rewrite.ScopedValue.create_name_values([target.value], [target.scope])
+        outcast_node = ms.rewrite.Node.create_call_function(_amp_cast_op, targets=outcast_targets, args=outcast_args)
+        stree.insert(stree.after(node), outcast_node)
+
+
+def _insert_cast_for_operators(stree, dtype, force_cast, *, white_list=None, black_list=None):
+    """insert cast for operators not in black_list."""
+    # get all nodes of stree exclude nodes in subtree.
+    all_nodes = stree.all_nodes(False)
+    for node in all_nodes:
+        if not node.get_targets():
+            continue
+        if _operator_need_cast(node, force_cast, white_list, black_list):
+            _insert_cast_for_operator(node, dtype)
+        elif node.get_node_type() == ms.rewrite.NodeType.Tree:
+            force_cast_ = force_cast or _net_need_cast(node, force_cast, white_list, black_list)
+            if not _precision_set_by_user(node.get_instance()):
+                subtree = node.get_sub_tree()
+                _insert_cast_for_operators(subtree, dtype, force_cast_, white_list=white_list, black_list=black_list)
 
 
 def _need_removed_cast_pair(node, dtype):
     """check whether the cast pairs should be removed."""
-    dtype_str = "mindspore.bfloat16" if dtype == mstype.bfloat16 else "mindspore.float16"
-    cast_dtypes = ms.rewrite.ScopedValue.create_name_values([dtype_str, "mindspore.float32"])
+    dtype_str = "bfloat16" if dtype == mstype.bfloat16 else "float16"
+    cast_dtypes = ms.rewrite.ScopedValue.create_name_values([dtype_str, "float32"], ["mindspore", "mindspore"])
     cast_dtype_f16 = cast_dtypes[0]
     cast_dtype_f32 = cast_dtypes[1]
-    # current node should be Cast Op to float32
+    # current node should be cast fp32
     if node.get_instance_type() != _amp_cast_op:
         return False
     node_cast_type = node.get_args()[1]
     if node_cast_type != cast_dtype_f32:
         return False
-    # all user nodes should be Cast Op to dtype or Cell with to_float(dtype)
+    # all user nodes should be cast fp16/bf16
     if not node.get_users():
         return False
     all_nodes = [ms.rewrite.Node(n) for n in node.get_handler().get_node_manager().nodes()]
     for user in node.get_users():
-        # If ControlFlow node(if statement) exists between current node and user node,
+        # If ControlFlow node(e.g. if, for, while) exists between current node and user node,
         # cast pair should not be removed.
         middle_nodes = all_nodes[all_nodes.index(node): all_nodes.index(user)]
         if any([n.get_node_type() == ms.rewrite.NodeType.ControlFlow for n in middle_nodes]):
             return False
-        if isinstance(user.get_instance(), nn.Cell):
-            to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-            if not (hasattr(user.get_instance(), to_float_flag) and getattr(user.get_instance(), to_float_flag)):
-                return False
-        elif user.get_instance_type() == _amp_cast_op:
-            user_cast_type = user.get_args()[1]
-            if user_cast_type != cast_dtype_f16:
-                return False
-        else:
+        if user.get_instance_type() != _amp_cast_op:
+            return False
+        user_cast_type = user.get_args()[1]
+        if user_cast_type != cast_dtype_f16:
             return False
+        # cast pair detected, check next user
+        continue
     return True
 
 
-def _removed_cast_pair_process(cast_f32_node):
-    """remove the duplicated cast operators."""
-    stree = cast_f32_node.get_symbol_tree()
-    cast_f32_users = cast_f32_node.get_users()
-    # remove cast f16 nodes
-    for user_node in cast_f32_users:
-        if user_node.get_instance_type() == _amp_cast_op:
-            cast_f16_node = user_node
-            # modify arguments using cast_f16's target[0] to cast_f32's args[0], which is f16 type
-            for cast_f16_user in cast_f16_node.get_users():
-                for idx, arg in enumerate(cast_f16_user.get_args()):
-                    if arg == cast_f16_node.get_targets()[0]:
-                        cast_f16_user.set_arg(idx, cast_f32_node.get_args()[0])
-            stree.erase(cast_f16_node)
-        # update args of cell f16 nodes
-        elif isinstance(user_node.get_instance(), nn.Cell):
-            cell_f16_node = user_node
-            for idx, arg in enumerate(cell_f16_node.get_args()):
-                if arg == cast_f32_node.get_targets()[0]:
-                    cell_f16_node.set_arg(idx, cast_f32_node.get_args()[0])
-    # remove the cast f32 node
-    stree.erase(cast_f32_node)
-
-
 def _remove_duplicated_cast(stree, dtype):
     """remove the duplicated cast operators."""
-    node_list = []
-    node_list.extend(list(stree.nodes()))
-    while node_list:
-        node = node_list.pop()
-        if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
-            for n in node.get_handler().node_list:
-                if n.get_node_type() == ms.rewrite.NodeType.Tree:
-                    _remove_duplicated_cast(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)), dtype)
-        elif node.get_node_type() == ms.rewrite.NodeType.Tree:
-            substree = ms.rewrite.TreeNodeHelper.get_sub_tree(node)
-            _remove_duplicated_cast(substree, dtype)
-        elif node.get_node_type() in [ms.rewrite.NodeType.CallFunction, ms.rewrite.NodeType.ControlFlow]:
-            if isinstance(node.get_handler(), ms.rewrite.node.NodeManager):
-                nodes = [ms.rewrite.Node(n) for n in node.get_handler().nodes()]
-                node_list.extend(nodes)
-        elif _need_removed_cast_pair(node, dtype):
-            _removed_cast_pair_process(node)
-
-
-def _auto_white_list(network, white_list, dtype):
-    """process the white list of network."""
-    stree = ms.rewrite.SymbolTree.create(network)
-    _insert_cast_operator_white_list(stree, white_list, dtype)
-    _remove_duplicated_cast(stree, dtype)
-    return stree.get_network()
-
-
-def _insert_cast_operator_black_list(stree, black_list, dtype):
-    """insert cast for operators not in black_list."""
-    allowed_list = []
-    # Ignore if net called ".to_float(dtype)"
-    net = stree.get_handler().get_origin_network()
-    to_float_flag = "bf16" if dtype == mstype.bfloat16 else "fp16"
-    if isinstance(net, nn.Cell) and hasattr(net, to_float_flag) and getattr(net, to_float_flag):
-        return
-    for node in stree.nodes(all_nodes=True):
-        if node.get_targets() is None:
-            continue
-        if node.get_node_type() == ms.rewrite.NodeType.CellContainer:
-            _insert_cast_for_cell_container(node, dtype, allowed_list, black_list=black_list)
-        elif isinstance(node.get_handler().get_node_manager(), ms.rewrite.node.CellContainer):
-            # nodes in CellContainer are processed by _insert_cast_for_cell_container
-            continue
-        elif node.get_instance_type() not in black_list and _allow_mix_precision(node, allowed_list, dtype):
-            _insert_cast_operator_process(node, dtype)
-
-
-def _remove_duplicated_cast_rewrite(stree, dtype):
-    """remove the duplicated cast operators."""
-    for node in stree.nodes(all_nodes=True):
+    all_nodes = list(stree.nodes(all_nodes=True))
+    for node in all_nodes:
         if _need_removed_cast_pair(node, dtype):
-            user_nodes = node.get_users()
-            # remove cast f16 nodes
-            for user_node in user_nodes:
-                if user_node.get_instance_type() == _amp_cast_op:
-                    stree.erase(user_node)
-            # remove the cast f32 node
+            incast_nodes = node.get_users()
+            # remove cast fp16/bf16 nodes
+            for incast_node in incast_nodes:
+                # get_target_users() return {target0: [(user0, arg_idx), ...], ...}
+                target_users = list(incast_node.get_target_users().values())
+                if not target_users or not target_users[0]:
+                    continue
+                for user_node, arg_idx in target_users[0]:
+                    user_node.set_arg(arg_idx, incast_node.get_args()[0])
+                stree.erase(incast_node)
+            # remove the cast fp32 node
             stree.erase(node)
 
 
-def _auto_black_list_rewrite(network, black_list, dtype):
+def _auto_mixed_precision_rewrite(network, dtype, *, white_list=None, black_list=None):
+    """Implement auto mixed precision by rewrite"""
+    if (white_list is None and black_list is None) or (white_list is not None and black_list is not None):
+        raise ValueError("For _auto_mixed_precision_rewrite, one of white_list and black_list must be provided.")
+    # enable rewrite configs for amp
+    ms.rewrite.common.namespace._ms_cells_to_subtree = True
+    ms.rewrite.parsers.assign_parser.AssignParser._share_one_implementation = True
+    # insert casts by rewrite
     stree = ms.rewrite.SymbolTree.create(network)
-    _insert_cast_operator_black_list(stree, black_list, dtype)
-    _remove_duplicated_cast_rewrite(stree, dtype)
-    return stree.get_network()
+    _insert_cast_for_operators(stree, dtype, False, white_list=white_list, black_list=black_list)
+    _remove_duplicated_cast(stree, dtype)
+    new_net = stree.get_network()
+    # disable rewrite configs
+    ms.rewrite.parsers.assign_parser.AssignParser._share_one_implementation = False
+    ms.rewrite.common.namespace._ms_cells_to_subtree = False
+    ms.rewrite.common.config.clear_caches()
+    return new_net
 
 
 def _auto_black_list(network, black_list, dtype):
@@ -410,7 +331,7 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
     :class:`mindspore.nn.LayerNorm`]
 
     For details on automatic mixed precision, refer to
-    `Automatic Mix Precision <https://www.mindspore.cn/tutorials/en/r2.2/advanced/mixed_precision.html>`_ .
+    `Automatic Mix Precision <https://www.mindspore.cn/tutorials/en/master/advanced/mixed_precision.html>`_ .
 
     Note:
         - Repeatedly calling mixed-precision interfaces, such as `custom_mixed_precision` and `auto_mixed_precision`,
@@ -441,7 +362,7 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
     Examples:
         >>> from mindspore import amp
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> network = LeNet5()
         >>> amp_level = "O1"
         >>> net = amp.auto_mixed_precision(network, amp_level)
@@ -462,16 +383,16 @@ def auto_mixed_precision(network, amp_level="O0", dtype=mstype.float16):
                        f"degradation.")
 
     if amp_level == "O1":
-        network = _auto_white_list(network, AMP_WHITE_LIST, dtype)
+        network = _auto_mixed_precision_rewrite(network, dtype, white_list=AMP_WHITE_LIST)
     elif amp_level == "O2":
         if MS_AMP_BY_REWRITE:
-            network = _auto_black_list_rewrite(network, AMP_BLACK_LIST, dtype)
+            network = _auto_mixed_precision_rewrite(network, dtype, black_list=AMP_BLACK_LIST)
         else:
             network = _auto_black_list(network, AMP_BLACK_LIST, dtype)
             network = _OutputTo32(network)
     elif amp_level == "O3":
         if MS_AMP_BY_REWRITE:
-            network = _auto_black_list_rewrite(network, [], dtype)
+            network = _auto_mixed_precision_rewrite(network, dtype, black_list=[])
         else:
             network.to_float(dtype)
             network = _OutputTo32(network)
@@ -638,8 +559,8 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
             - 'O1': Cast the operators in white_list to float16, the remaining operators are kept in float32.
               The operators in the whitelist: [Conv1d, Conv2d, Conv3d, Conv1dTranspose, Conv2dTranspose,
               Conv3dTranspose, Dense, LSTMCell, RNNCell, GRUCell, MatMul, BatchMatMul, PReLU, ReLU, Ger].
-            - 'O2': Cast network to float16, keep batchnorm and `loss_fn` (if set) run in float32,
-              using dynamic loss scale.
+            - 'O2': Cast network to float16, keep `mindspore.nn.BatchNorm` series interface,
+              :class:`mindspore.nn.LayerNorm` and `loss_fn` (if set) run in float32, using dynamic loss scale.
             - 'O3': Cast network to float16, with additional property `keep_batchnorm_fp32=False` .
             - 'auto': Set to level to recommended level in different devices. Set level to 'O2' on GPU, Set
               level to 'O3' Ascend. The recommended level is chosen by the export experience, not applicable to all
@@ -670,13 +591,13 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
             take no effect on this property.
 
     Raises:
-        ValueError: If device is CPU, property `loss_scale_manager` is not `None` or `FixedLossScaleManager`
-            (with property `drop_overflow_update=False` ).
+        ValueError: If device is CPU, property `loss_scale_manager` is not `None` or
+            :class:`mindspore.amp.FixedLossScaleManager` (with property `drop_overflow_update=False` ).
 
     Examples:
         >>> from mindspore import amp, nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> network = LeNet5()
         >>> net_loss = nn.SoftmaxCrossEntropyWithLogits(reduction="mean")
         >>> net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)
@@ -823,7 +744,7 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=m
     Examples:
         >>> from mindspore import amp, nn
         >>> # Define the network structure of LeNet5. Refer to
-        >>> # https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
+        >>> # https://gitee.com/mindspore/docs/blob/master/docs/mindspore/code/lenet.py
         >>> net = LeNet5()
         >>> custom_white_list = amp.get_white_list()
         >>> custom_white_list.append(nn.Flatten)
@@ -844,11 +765,11 @@ def custom_mixed_precision(network, *, white_list=None, black_list=None, dtype=m
 
     if white_list is not None:
         _list_check(white_list, "white_list")
-        network = _auto_white_list(network, white_list, dtype)
+        network = _auto_mixed_precision_rewrite(network, dtype, white_list=white_list)
     else:
         _list_check(black_list, "black_list")
         if MS_AMP_BY_REWRITE:
-            network = _auto_black_list_rewrite(network, black_list, dtype)
+            network = _auto_mixed_precision_rewrite(network, dtype, black_list=black_list)
         else:
             network = _auto_black_list(network, black_list, dtype)
             network = _OutputTo32(network)
@@ -883,7 +804,8 @@ def _list_check(custom_list: list, list_name: str):
             if elem not in custom_list:
                 logger.warning(f"{elem} is removed from internal black list.")
 
-def _config_amp(*, enable_rewrite: bool = None, cast_op: type = None): # pylint: disable=unused-variable
+
+def _config_amp(*, enable_rewrite: bool = None, cast_op: types.FunctionType = None): # pylint: disable=unused-variable
     """Configure auto mixed precision."""
     global MS_AMP_BY_REWRITE
     global _amp_cast_op

mindspore/train/anf_ir_pb2.py

@@ -20,7 +20,7 @@ DESCRIPTOR = _descriptor.FileDescriptor(
   syntax='proto2',
   serialized_options=None,
   create_key=_descriptor._internal_create_key,
-  serialized_pb=b'\n\x0c\x61nf_ir.proto\x12\x0emindspore.irpb\"\xdb\x04\n\nValueProto\x12\'\n\x05\x64type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x10\n\x08\x62ool_val\x18\x02 \x01(\x08\x12\x0f\n\x07int_val\x18\x03 \x01(\x03\x12\x10\n\x08uint_val\x18\x04 \x01(\x04\x12\x11\n\tfloat_val\x18\x05 \x01(\x02\x12\x12\n\ndouble_val\x18\x06 \x01(\x01\x12\x0f\n\x07str_val\x18\x07 \x01(\t\x12/\n\ntensor_val\x18\x08 \x01(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12)\n\x05graph\x18\t \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12\x11\n\tbool_vals\x18\n \x03(\x08\x12\x10\n\x08int_vals\x18\x0b \x03(\x03\x12\x11\n\tuint_vals\x18\x0c \x03(\x04\x12\x12\n\nfloat_vals\x18\r \x03(\x02\x12\x13\n\x0b\x64ouble_vals\x18\x0e \x03(\x01\x12\x10\n\x08str_vals\x18\x0f \x03(\t\x12\x30\n\x0btensor_vals\x18\x10 \x03(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12*\n\x06graphs\x18\x11 \x03(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12*\n\x06values\x18\x12 \x03(\x0b\x32\x1a.mindspore.irpb.ValueProto\x12\x31\n\x08\x64ict_val\x18\x13 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\x12+\n\x08type_val\x18\x14 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"I\n\x0e\x41ttributeProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"I\n\x0fNamedValueProto\x12\x0b\n\x03key\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"t\n\x10TensorShapeProto\x12\x37\n\x03\x64im\x18\x01 \x03(\x0b\x32*.mindspore.irpb.TensorShapeProto.Dimension\x1a\'\n\tDimension\x12\x0c\n\x04size\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\"\xda\x02\n\tTypeProto\x12+\n\tdata_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x37\n\x0btensor_type\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TypeProto.TensorH\x00\x12;\n\rsequence_type\x18\x03 \x01(\x0b\x32\".mindspore.irpb.TypeProto.SequenceH\x00\x1a\x66\n\x06Tensor\x12+\n\telem_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12/\n\x05shape\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TensorShapeProto\x1a\x39\n\x08Sequence\x12-\n\nelem_types\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.TypeProtoB\x07\n\x05value\"x\n\x0eParameterProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12/\n\x0b\x64\x65\x66\x61ult_val\x18\x03 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"D\n\x0bOutputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"z\n\nInputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x31\n\x04type\x18\x02 \x01(\x0e\x32#.mindspore.irpb.InputProto.EdgeType\"+\n\x08\x45\x64geType\x12\r\n\tDATA_EDGE\x10\x00\x12\x10\n\x0c\x43ONTROL_EDGE\x10\x01\"\x83\x02\n\tNodeProto\x12)\n\x05input\x18\x01 \x03(\x0b\x32\x1a.mindspore.irpb.InputProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0f\n\x07op_type\x18\x03 \x01(\t\x12\r\n\x05scope\x18\x04 \x01(\t\x12\x31\n\tattribute\x18\x05 \x03(\x0b\x32\x1e.mindspore.irpb.AttributeProto\x12.\n\x0boutput_type\x18\x06 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12\x10\n\x08output_i\x18\x07 \x01(\x04\x12\x11\n\tfull_name\x18\x08 \x01(\t\x12\x15\n\rinstance_name\x18\n \x01(\t\"\xb0\x01\n\nModelProto\x12\x12\n\nir_version\x18\x01 \x01(\x03\x12\x0e\n\x06\x64omain\x18\x02 \x01(\t\x12\x15\n\rmodel_version\x18\x03 \x01(\x03\x12)\n\x05graph\x18\x04 \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12<\n\x12metadata_operators\x18\x05 \x01(\x0b\x32 .mindspore.irpb.OperatorSetProto\"?\n\rOperatorProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06\x63onfig\x18\x02 \x01(\x0c\x12\x10\n\x08obj_info\x18\x03 \x01(\x0c\"U\n\x10OperatorSetProto\x12\x30\n\toperators\x18\x01 \x03(\x0b\x32\x1d.mindspore.irpb.OperatorProto\x12\x0f\n\x07summary\x18\x02 \x01(\t\"\xda\x01\n\nGraphProto\x12\'\n\x04node\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.NodeProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x32\n\nparameters\x18\x03 \x03(\x0b\x32\x1e.mindspore.irpb.ParameterProto\x12,\n\x07outputs\x18\x04 \x03(\x0b\x32\x1b.mindspore.irpb.OutputProto\x12\x33\n\nconst_vals\x18\x05 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\"\xd4\x01\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12+\n\tdata_type\x18\x02 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x16\n\nfloat_data\x18\x03 \x03(\x02\x42\x02\x10\x01\x12\x16\n\nint32_data\x18\x04 \x03(\x05\x42\x02\x10\x01\x12\x16\n\nint64_data\x18\x05 \x03(\x03\x42\x02\x10\x01\x12\x17\n\x0b\x64ouble_data\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x17\n\x0buint64_data\x18\x07 \x03(\x04\x42\x02\x10\x01\x12\x10\n\x08raw_data\x18\x08 \x01(\x0c*/\n\x07Version\x12\x14\n\x10UNKNOWWN_VERSION\x10\x00\x12\x0e\n\nIR_VERSION\x10\x01*\xee\x05\n\x08\x44\x61taType\x12\x10\n\x0c\x44T_UNDEFINED\x10\x00\x12\x0b\n\x07\x44T_BOOL\x10\x01\x12\x0b\n\x07\x44T_INT8\x10\x02\x12\x0c\n\x08\x44T_INT16\x10\x03\x12\x0c\n\x08\x44T_INT32\x10\x04\x12\x0c\n\x08\x44T_INT64\x10\x05\x12\x0c\n\x08\x44T_UINT8\x10\x06\x12\r\n\tDT_UINT16\x10\x07\x12\r\n\tDT_UINT32\x10\x08\x12\r\n\tDT_UINT64\x10\t\x12\x0e\n\nDT_FLOAT16\x10\n\x12\x0e\n\nDT_FLOAT32\x10\x0b\x12\x0e\n\nDT_FLOAT64\x10\x0c\x12\r\n\tDT_STRING\x10\r\x12\r\n\tDT_TENSOR\x10\x0e\x12\x0c\n\x08\x44T_GRAPH\x10\x0f\x12\x0c\n\x08\x44T_BOOLS\x10\x10\x12\x0c\n\x08\x44T_INTS8\x10\x11\x12\r\n\tDT_INTS16\x10\x12\x12\r\n\tDT_INTS32\x10\x13\x12\r\n\tDT_INTS64\x10\x14\x12\r\n\tDT_UINTS8\x10\x15\x12\x0e\n\nDT_UINTS16\x10\x16\x12\x0e\n\nDT_UINTS32\x10\x17\x12\x0e\n\nDT_UINTS64\x10\x18\x12\x0f\n\x0b\x44T_FLOATS16\x10\x19\x12\x0f\n\x0b\x44T_FLOATS32\x10\x1a\x12\x0f\n\x0b\x44T_FLOATS64\x10\x1b\x12\x0e\n\nDT_STRINGS\x10\x1c\x12\x0e\n\nDT_TENSORS\x10\x1d\x12\r\n\tDT_GRAPHS\x10\x1e\x12\x0c\n\x08\x44T_TUPLE\x10\x1f\x12\x0b\n\x07\x44T_LIST\x10 \x12\x0b\n\x07\x44T_DICT\x10!\x12\x0b\n\x07\x44T_NONE\x10\"\x12\x0f\n\x0b\x44T_SYM_INST\x10#\x12\x0f\n\x0b\x44T_BASE_INT\x10$\x12\x10\n\x0c\x44T_BASE_UINT\x10%\x12\x11\n\rDT_BASE_FLOAT\x10&\x12\x0b\n\x07\x44T_TYPE\x10\'\x12\n\n\x06\x44T_ANY\x10(\x12\r\n\tDT_REFKEY\x10)\x12\n\n\x06\x44T_REF\x10*\x12\x10\n\x0c\x44T_COMPLEX64\x10+\x12\x11\n\rDT_COMPLEX128\x10,\x12\x13\n\x0f\x44T_BASE_COMPLEX\x10-\x12\x0f\n\x0b\x44T_BFLOAT16\x10.\x12\x10\n\x0c\x44T_BFLOATS16\x10/'
+  serialized_pb=b'\n\x0c\x61nf_ir.proto\x12\x0emindspore.irpb\"\xdb\x04\n\nValueProto\x12\'\n\x05\x64type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x10\n\x08\x62ool_val\x18\x02 \x01(\x08\x12\x0f\n\x07int_val\x18\x03 \x01(\x03\x12\x10\n\x08uint_val\x18\x04 \x01(\x04\x12\x11\n\tfloat_val\x18\x05 \x01(\x02\x12\x12\n\ndouble_val\x18\x06 \x01(\x01\x12\x0f\n\x07str_val\x18\x07 \x01(\t\x12/\n\ntensor_val\x18\x08 \x01(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12)\n\x05graph\x18\t \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12\x11\n\tbool_vals\x18\n \x03(\x08\x12\x10\n\x08int_vals\x18\x0b \x03(\x03\x12\x11\n\tuint_vals\x18\x0c \x03(\x04\x12\x12\n\nfloat_vals\x18\r \x03(\x02\x12\x13\n\x0b\x64ouble_vals\x18\x0e \x03(\x01\x12\x10\n\x08str_vals\x18\x0f \x03(\t\x12\x30\n\x0btensor_vals\x18\x10 \x03(\x0b\x32\x1b.mindspore.irpb.TensorProto\x12*\n\x06graphs\x18\x11 \x03(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12*\n\x06values\x18\x12 \x03(\x0b\x32\x1a.mindspore.irpb.ValueProto\x12\x31\n\x08\x64ict_val\x18\x13 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\x12+\n\x08type_val\x18\x14 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"I\n\x0e\x41ttributeProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"I\n\x0fNamedValueProto\x12\x0b\n\x03key\x18\x01 \x01(\t\x12)\n\x05value\x18\x02 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"t\n\x10TensorShapeProto\x12\x37\n\x03\x64im\x18\x01 \x03(\x0b\x32*.mindspore.irpb.TensorShapeProto.Dimension\x1a\'\n\tDimension\x12\x0c\n\x04size\x18\x01 \x01(\x03\x12\x0c\n\x04name\x18\x02 \x01(\t\"\xda\x02\n\tTypeProto\x12+\n\tdata_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x37\n\x0btensor_type\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TypeProto.TensorH\x00\x12;\n\rsequence_type\x18\x03 \x01(\x0b\x32\".mindspore.irpb.TypeProto.SequenceH\x00\x1a\x66\n\x06Tensor\x12+\n\telem_type\x18\x01 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12/\n\x05shape\x18\x02 \x01(\x0b\x32 .mindspore.irpb.TensorShapeProto\x1a\x39\n\x08Sequence\x12-\n\nelem_types\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.TypeProtoB\x07\n\x05value\"x\n\x0eParameterProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12/\n\x0b\x64\x65\x66\x61ult_val\x18\x03 \x01(\x0b\x32\x1a.mindspore.irpb.ValueProto\"D\n\x0bOutputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\'\n\x04type\x18\x02 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\"z\n\nInputProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x31\n\x04type\x18\x02 \x01(\x0e\x32#.mindspore.irpb.InputProto.EdgeType\"+\n\x08\x45\x64geType\x12\r\n\tDATA_EDGE\x10\x00\x12\x10\n\x0c\x43ONTROL_EDGE\x10\x01\"\x83\x02\n\tNodeProto\x12)\n\x05input\x18\x01 \x03(\x0b\x32\x1a.mindspore.irpb.InputProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0f\n\x07op_type\x18\x03 \x01(\t\x12\r\n\x05scope\x18\x04 \x01(\t\x12\x31\n\tattribute\x18\x05 \x03(\x0b\x32\x1e.mindspore.irpb.AttributeProto\x12.\n\x0boutput_type\x18\x06 \x01(\x0b\x32\x19.mindspore.irpb.TypeProto\x12\x10\n\x08output_i\x18\x07 \x01(\x04\x12\x11\n\tfull_name\x18\x08 \x01(\t\x12\x15\n\rinstance_name\x18\n \x01(\t\"\xb0\x01\n\nModelProto\x12\x12\n\nir_version\x18\x01 \x01(\x03\x12\x0e\n\x06\x64omain\x18\x02 \x01(\t\x12\x15\n\rmodel_version\x18\x03 \x01(\x03\x12)\n\x05graph\x18\x04 \x01(\x0b\x32\x1a.mindspore.irpb.GraphProto\x12<\n\x12metadata_operators\x18\x05 \x01(\x0b\x32 .mindspore.irpb.OperatorSetProto\"?\n\rOperatorProto\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06\x63onfig\x18\x02 \x01(\x0c\x12\x10\n\x08obj_info\x18\x03 \x01(\x0c\"U\n\x10OperatorSetProto\x12\x30\n\toperators\x18\x01 \x03(\x0b\x32\x1d.mindspore.irpb.OperatorProto\x12\x0f\n\x07summary\x18\x02 \x01(\t\"\xda\x01\n\nGraphProto\x12\'\n\x04node\x18\x01 \x03(\x0b\x32\x19.mindspore.irpb.NodeProto\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x32\n\nparameters\x18\x03 \x03(\x0b\x32\x1e.mindspore.irpb.ParameterProto\x12,\n\x07outputs\x18\x04 \x03(\x0b\x32\x1b.mindspore.irpb.OutputProto\x12\x33\n\nconst_vals\x18\x05 \x03(\x0b\x32\x1f.mindspore.irpb.NamedValueProto\"\xd4\x01\n\x0bTensorProto\x12\x0c\n\x04\x64ims\x18\x01 \x03(\x03\x12+\n\tdata_type\x18\x02 \x01(\x0e\x32\x18.mindspore.irpb.DataType\x12\x16\n\nfloat_data\x18\x03 \x03(\x02\x42\x02\x10\x01\x12\x16\n\nint32_data\x18\x04 \x03(\x05\x42\x02\x10\x01\x12\x16\n\nint64_data\x18\x05 \x03(\x03\x42\x02\x10\x01\x12\x17\n\x0b\x64ouble_data\x18\x06 \x03(\x01\x42\x02\x10\x01\x12\x17\n\x0buint64_data\x18\x07 \x03(\x04\x42\x02\x10\x01\x12\x10\n\x08raw_data\x18\x08 \x01(\x0c*/\n\x07Version\x12\x14\n\x10UNKNOWWN_VERSION\x10\x00\x12\x0e\n\nIR_VERSION\x10\x01*\x89\x06\n\x08\x44\x61taType\x12\x10\n\x0c\x44T_UNDEFINED\x10\x00\x12\x0b\n\x07\x44T_BOOL\x10\x01\x12\x0b\n\x07\x44T_INT8\x10\x02\x12\x0c\n\x08\x44T_INT16\x10\x03\x12\x0c\n\x08\x44T_INT32\x10\x04\x12\x0c\n\x08\x44T_INT64\x10\x05\x12\x0c\n\x08\x44T_UINT8\x10\x06\x12\r\n\tDT_UINT16\x10\x07\x12\r\n\tDT_UINT32\x10\x08\x12\r\n\tDT_UINT64\x10\t\x12\x0e\n\nDT_FLOAT16\x10\n\x12\x0e\n\nDT_FLOAT32\x10\x0b\x12\x0e\n\nDT_FLOAT64\x10\x0c\x12\r\n\tDT_STRING\x10\r\x12\r\n\tDT_TENSOR\x10\x0e\x12\x0c\n\x08\x44T_GRAPH\x10\x0f\x12\x0c\n\x08\x44T_BOOLS\x10\x10\x12\x0c\n\x08\x44T_INTS8\x10\x11\x12\r\n\tDT_INTS16\x10\x12\x12\r\n\tDT_INTS32\x10\x13\x12\r\n\tDT_INTS64\x10\x14\x12\r\n\tDT_UINTS8\x10\x15\x12\x0e\n\nDT_UINTS16\x10\x16\x12\x0e\n\nDT_UINTS32\x10\x17\x12\x0e\n\nDT_UINTS64\x10\x18\x12\x0f\n\x0b\x44T_FLOATS16\x10\x19\x12\x0f\n\x0b\x44T_FLOATS32\x10\x1a\x12\x0f\n\x0b\x44T_FLOATS64\x10\x1b\x12\x0e\n\nDT_STRINGS\x10\x1c\x12\x0e\n\nDT_TENSORS\x10\x1d\x12\r\n\tDT_GRAPHS\x10\x1e\x12\x0c\n\x08\x44T_TUPLE\x10\x1f\x12\x0b\n\x07\x44T_LIST\x10 \x12\x0b\n\x07\x44T_DICT\x10!\x12\x0b\n\x07\x44T_NONE\x10\"\x12\x0f\n\x0b\x44T_SYM_INST\x10#\x12\x0f\n\x0b\x44T_BASE_INT\x10$\x12\x10\n\x0c\x44T_BASE_UINT\x10%\x12\x11\n\rDT_BASE_FLOAT\x10&\x12\x0b\n\x07\x44T_TYPE\x10\'\x12\n\n\x06\x44T_ANY\x10(\x12\r\n\tDT_REFKEY\x10)\x12\n\n\x06\x44T_REF\x10*\x12\x10\n\x0c\x44T_COMPLEX64\x10+\x12\x11\n\rDT_COMPLEX128\x10,\x12\x13\n\x0f\x44T_BASE_COMPLEX\x10-\x12\x0f\n\x0b\x44T_BFLOAT16\x10.\x12\x10\n\x0c\x44T_BFLOATS16\x10/\x12\x0b\n\x07\x44T_INT4\x10\x30\x12\x0c\n\x08\x44T_SLICE\x10\x31'
 )
 
 _VERSION = _descriptor.EnumDescriptor(
@@ -296,11 +296,21 @@ _DATATYPE = _descriptor.EnumDescriptor(
       serialized_options=None,
       type=None,
       create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='DT_INT4', index=48, number=48,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
+    _descriptor.EnumValueDescriptor(
+      name='DT_SLICE', index=49, number=49,
+      serialized_options=None,
+      type=None,
+      create_key=_descriptor._internal_create_key),
   ],
   containing_type=None,
   serialized_options=None,
   serialized_start=2650,
-  serialized_end=3400,
+  serialized_end=3427,
 )
 _sym_db.RegisterEnumDescriptor(_DATATYPE)
 
@@ -355,6 +365,8 @@ DT_COMPLEX128 = 44
 DT_BASE_COMPLEX = 45
 DT_BFLOAT16 = 46
 DT_BFLOATS16 = 47
+DT_INT4 = 48
+DT_SLICE = 49
 
 
 _INPUTPROTO_EDGETYPE = _descriptor.EnumDescriptor(

mindspore/train/callback/__init__.py

@@ -35,7 +35,10 @@ from mindspore.train.callback._early_stop import EarlyStopping
 from mindspore.train.callback._reduce_lr_on_plateau import ReduceLROnPlateau
 from mindspore.train.callback._on_request_exit import OnRequestExit
 from mindspore.train.callback._backup_and_restore import BackupAndRestore
+from mindspore.train.callback._flops_collector import FlopsUtilizationCollector
+from mindspore.train.callback._mindio_ttp import MindIOTTPAdapter
 
-__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint",
+__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint", "FlopsUtilizationCollector",
            "SummaryCollector", "CheckpointConfig", "RunContext", "LearningRateScheduler", "SummaryLandscape",
-           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore"]
+           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore",
+           "MindIOTTPAdapter"]