mindspore 1.10.0__cp37-cp37m-win_amd64.whl → 2.0.0rc1__cp37-cp37m-win_amd64.whl

mindspore/train/__init__.py

@@ -26,19 +26,21 @@ from mindspore.train.amp import build_train_network
 from mindspore.train.loss_scale_manager import LossScaleManager, FixedLossScaleManager, DynamicLossScaleManager
 from mindspore.train.serialization import save_checkpoint, load_checkpoint, load_param_into_net, export, \
     load, parse_print, build_searched_strategy, merge_sliced_parameter, load_distributed_checkpoint, \
-    async_ckpt_thread_status, restore_group_info_list, convert_model
+    async_ckpt_thread_status, restore_group_info_list, convert_model, obfuscate_model
 from mindspore.train.callback import Callback, LossMonitor, TimeMonitor, ModelCheckpoint, SummaryCollector, \
     CheckpointConfig, RunContext, LearningRateScheduler, SummaryLandscape, \
-    History, LambdaCallback, ReduceLROnPlateau, EarlyStopping
+    History, LambdaCallback, ReduceLROnPlateau, EarlyStopping, OnRequestExit, BackupAndRestore
 from mindspore.train.summary import SummaryRecord
 from mindspore.train.train_thor import ConvertNetUtils, ConvertModelUtils
+from mindspore.train.metrics import *
 from mindspore.train.data_sink import data_sink
 
-__all__ = ["Model", "DatasetHelper", "amp", "connect_network_with_dataset", "build_train_network", "LossScaleManager",
+__all__ = ["Model", "DatasetHelper", "connect_network_with_dataset", "build_train_network", "LossScaleManager",
            "FixedLossScaleManager", "DynamicLossScaleManager", "save_checkpoint", "load_checkpoint",
            "load_param_into_net", "export", "load", "parse_print", "build_searched_strategy", "merge_sliced_parameter",
            "load_distributed_checkpoint", "async_ckpt_thread_status", "restore_group_info_list", "convert_model",
-           "data_sink"]
+           "data_sink", "obfuscate_model"]
 __all__.extend(callback.__all__)
 __all__.extend(summary.__all__)
 __all__.extend(train_thor.__all__)
+__all__.extend(metrics.__all__)

mindspore/train/_utils.py

@@ -20,16 +20,18 @@ from collections.abc import Iterable
 import numpy as np
 
 from mindspore.common.tensor import Tensor
+from mindspore._c_expression import Tensor as Tensor_
 from mindspore.common.dtype import dtype_to_nptype, pytype_to_dtype
 from mindspore.common import dtype as mstype
 from mindspore import log as logger
-from mindspore._checkparam import Validator
+from mindspore import _checkparam as Validator
 from mindspore.common.api import _cell_graph_executor
 from mindspore.train.mind_ir_pb2 import ModelProto as mindir_model
 from mindspore.train.checkpoint_pb2 import Checkpoint
 from mindspore.train.node_strategy_pb2 import ParallelStrategyMap as ckpt_strategy
 from mindspore.train.lineage_pb2 import DatasetGraph, TrainLineage, EvaluationLineage, UserDefinedInfo
 from mindspore.parallel._parallel_serialization import _make_dir
+from mindspore.ops.operations import debug_ops
 
 
 def _convert_type(types):
@@ -135,27 +137,48 @@ def _construct_input_tensors(dataset_types, dataset_shapes, device_number=1):
     return tensor_list_run, tensor_list_compile
 
 
-def _check_to_numpy(plugin, tensor):
+def _check_to_numpy(plugin, tensor, prim=None):
     """Check the tensor and return a numpy.ndarray."""
     np_value = tensor.asnumpy()
     np_value = np_value.copy()
+    summary_name = plugin.capitalize() + "Summary" if prim else "SummaryRecord"
     if plugin == 'scalar':
         if np_value.size == 1:
             return np_value
-        raise ValueError('The tensor holds more than one value, but the scalar plugin expects on value.')
+        raise ValueError(
+            f'For "{summary_name}", the v rank must be less than or equal to 1, but got {len(np_value)}.')
     if plugin == 'image':
         if np_value.ndim == 4:
             return np_value
-        raise ValueError('The tensor seems not to hold a valid image.')
+        raise ValueError(f'For "{summary_name}", The tensor seems not to hold a valid image.')
     if plugin in ('tensor', 'histogram'):
         if np_value.ndim > 0:
             return np_value
-        raise ValueError('The tensor should not be empty.')
+        raise ValueError(f'For "{summary_name}", The value should not be empty.')
     return np_value
 
 
+def check_summary_param(summary_name, tag, tensor):
+    """Checks the tag is valid for summary."""
+    plugin = summary_name.split('Summary')[0].lower()
+    try:
+        if not isinstance(tag, str) or not tag:
+            raise TypeError(f'For "{summary_name}", the name must be valid string, but got "{tag}".')
+        if not isinstance(tensor, (Tensor, Tensor_)):
+            raise TypeError(f'For "{summary_name}", the parameter "value" expect to be Tensor, '
+                            f'but got {type(tensor).__name__}')
+        _check_to_numpy(plugin, tensor, prim=True)
+    except TypeError as err:
+        raise TypeError(err)
+    except ValueError as err:
+        raise ValueError(err)
+    finally:
+        debug_ops.SUMMARY_TENSOR_CACHE = []
+
+
 def _check_lineage_value(plugin, value):
     """Check the lineage value."""
+
     def raises(plugin, prototype):
         raise TypeError(f'Plugin {repr(plugin)} expects a {prototype.__name__} value.')
 

mindspore/train/amp.py

@@ -15,9 +15,9 @@
 """Auto mixed precision."""
 from __future__ import absolute_import
 
+import mindspore as ms
 from mindspore import nn
-from mindspore._checkparam import Validator as validator
-from mindspore._checkparam import Rel
+from mindspore import _checkparam as validator
 from mindspore.common import dtype as mstype
 from mindspore.nn.wrap.cell_wrapper import _TrainPipelineAccuStepCell
 from mindspore.nn.wrap.loss_scale import _TrainPipelineWithLossScaleCell
@@ -25,73 +25,231 @@ from mindspore.ops import functional as F
 from mindspore.parallel._utils import _get_pipeline_stages
 from mindspore.train.loss_scale_manager import DynamicLossScaleManager, LossScaleManager
 from mindspore import boost, context
+from mindspore.ops import operations as P
+from mindspore.ops import Primitive
+from mindspore import log as logger
 
 
-AMP_WHITE_LIST = (
-    nn.Dense,
+STREE = None
+
+
+AMP_WHITE_LIST = [
     nn.Conv1d,
     nn.Conv2d,
     nn.Conv3d,
     nn.Conv1dTranspose,
     nn.Conv2dTranspose,
-    nn.Conv3dTranspose
-)
-
-AMP_BLACK_LIST = (
+    nn.Conv3dTranspose,
+    nn.Dense,
+    nn.LSTMCell,
+    nn.RNNCell,
+    nn.GRUCell,
+    P.Conv2D,
+    P.Conv3D,
+    P.Conv2DTranspose,
+    P.Conv3DTranspose,
+    P.Conv2DBackpropInput,
+    P.MatMul,
+    P.BatchMatMul,
+    P.PReLU,
+    P.ReLU,
+    P.Ger
+]
+
+
+AMP_BLACK_LIST = [
     nn.BatchNorm1d,
     nn.BatchNorm2d,
     nn.BatchNorm3d,
     nn.LayerNorm
-)
+]
 
 
 class _OutputTo16(nn.Cell):
     """Wrap cell for amp. Cast network output back to float16."""
-
-    def __init__(self, op):
+    def __init__(self, backbone):
         super(_OutputTo16, self).__init__(auto_prefix=False)
-        self._op = op
+        self._backbone = backbone
+        if isinstance(backbone, nn.Cell) and backbone.jit_config_dict:
+            self._jit_config_dict = backbone.jit_config_dict
 
     def construct(self, x):
-        return F.cast(self._op(x), mstype.float16)
+        return F.cast(self._backbone(x), mstype.float16)
 
 
 class _OutputTo32(nn.Cell):
-    "Wrap loss for amp. Cast network output back to float32"
-
+    """Wrap loss for amp. Cast network output back to float32."""
     def __init__(self, backbone):
         super(_OutputTo32, self).__init__(auto_prefix=False)
         self._backbone = backbone
-        self._jit_config_dict = backbone._jit_config_dict
+        if isinstance(backbone, nn.Cell) and backbone.jit_config_dict:
+            self._jit_config_dict = backbone.jit_config_dict
 
     def construct(self, *inputs):
         out = self._backbone(*inputs)
         return F.mixed_precision_cast(mstype.float32, out)
 
 
-def _auto_white_list(network, white_list=None):
-    """process the white list of network."""
-    if white_list is None:
-        white_list = AMP_WHITE_LIST
-    cells = network.name_cells()
-    change = False
-    for name in cells:
-        subcell = cells[name]
-        if subcell == network:
+def _allow_mix_precision(node, allowed_list) -> 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 P.Cast()
+        3) to_float(mindspore.float16) is not set in Cell
+    """
+    if node.get_instance() in allowed_list:
+        return True
+    if not issubclass(node.get_instance_type(), (Primitive, nn.Cell)):
+        return False
+    if isinstance(node.get_instance(), P.Cast):
+        return False
+    if issubclass(node.get_instance_type(), nn.Cell):
+        # if cell is already in allowed_list, it means to_float(mindspore.float16) is set by amp.
+        # if cell is not in allowed_list, but has to_float(mindspore.float16),
+        # it means to_float(mindspore.float16) is set by user.
+        if node.get_instance().to_float_fp16:
+            return False
+    allowed_list.append(node.get_instance())
+    return True
+
+
+def _insert_cast_operator_process(node, stree):
+    """insert cast for operators in white_list."""
+    new_cast_node = None
+    # insert cast float16 before the primitive operators
+    if issubclass(node.get_instance_type(), Primitive):
+        for idx in range(len(node.get_inputs())):
+            position = stree.before(node)
+            new_node = P.Cast()
+            arg = ms.rewrite.ScopedValue.create_name_values([node.get_inputs()[idx].get_targets()[0].value,
+                                                             "mindspore.float16"])
+            new_cast_node = ms.rewrite.Node.create_call_cell(new_node,
+                                                             targets=['x_cast_{}'.format(node.get_name())],
+                                                             args=arg,
+                                                             name='incast_{}{}'.format(node.get_name(), idx))
+            stree.insert(position, new_cast_node)
+            node.set_arg_by_node(idx, new_cast_node)
+    # insert cast float16 before the Cell operators
+    elif issubclass(node.get_instance_type(), nn.Cell):
+        node.get_instance().to_float(mstype.float16)
+    # ignore if subclass is not one of (Primitive, nn.Cell)
+    else:
+        return
+
+    # insert cast float32 after the operators
+    position = stree.after(node)
+    new_node = P.Cast()
+    arg = 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=['x_cast_{}'.format(node.get_name())],
+                                                     args=arg,
+                                                     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
-        if isinstance(subcell, white_list):
-            network._cells[name] = _OutputTo32(subcell.to_float(mstype.float16))
-            change = True
+        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):
+    """insert cast for operators in white_list."""
+    allowed_list = []
+    for node in stree.nodes():
+        if node.get_targets() is None:
+            continue
+        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:
+                    _insert_cast_operator_white_list(ms.rewrite.TreeNodeHelper.get_sub_tree(ms.rewrite.Node(n)),
+                                                     white_list)
+        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)
+        elif node.get_instance_type() in white_list and _allow_mix_precision(node, allowed_list):
+            _insert_cast_operator_process(node, stree)
+
+
+def _need_removed_cast_pair(node):
+    """check whether the cast pairs should be removed."""
+    cast_dtypes = ms.rewrite.ScopedValue.create_name_values(["mindspore.float16", "mindspore.float32"])
+    cast_dtype_f16 = cast_dtypes[0]
+    cast_dtype_f32 = cast_dtypes[1]
+    # current node should be P.Cast()(x, mindspore.float32)
+    if node.get_instance_type() != P.Cast:
+        return False
+    node_cast_type = node.get_args()[1]
+    if node_cast_type != cast_dtype_f32:
+        return False
+    # all user nodes should be P.Cast()(x, mindspore.float16) or Cell with to_float(mindspore.float16)
+    if not node.get_users():
+        return False
+    for user in node.get_users():
+        if isinstance(user.get_instance(), nn.Cell):
+            if not user.get_instance().to_float_fp16:
+                return False
+        elif user.get_instance_type() == P.Cast:
+            user_cast_type = user.get_args()[1]
+            if user_cast_type != cast_dtype_f16:
+                return False
         else:
-            _auto_white_list(subcell, white_list)
-    if isinstance(network, nn.SequentialCell) and change:
-        network.cell_list = list(network.cells())
+            return False
+    return True
+
+
+def _removed_cast_pair_process(stree, cast_f32_node):
+    """remove the duplicated cast operators."""
+    for user_node in cast_f32_node.get_users():
+        # remove cast f16 nodes
+        if user_node.get_instance_type() == P.Cast:
+            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_node(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_node(cast_f32_node)
+
+
+def _remove_duplicated_cast(stree):
+    """remove the duplicated cast operators."""
+    for node in stree.nodes():
+        if node.get_targets() is None:
+            continue
+        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)))
+        elif node.get_node_type() == ms.rewrite.NodeType.Tree:
+            substree = ms.rewrite.TreeNodeHelper.get_sub_tree(node)
+            _remove_duplicated_cast(substree)
+        elif _need_removed_cast_pair(node):
+            _removed_cast_pair_process(stree, node)
+
+
+def _auto_white_list(network, white_list):
+    """process the white list of network."""
+    global STREE
+    STREE = ms.rewrite.SymbolTree.create(network)
+    _insert_cast_operator_white_list(STREE, white_list)
+    _remove_duplicated_cast(STREE)
+    return STREE.get_network()
 
 
-def _auto_black_list(network, black_list=None):
+def _auto_black_list(network, black_list):
     """process the black list of network."""
-    if black_list is None:
-        black_list = AMP_BLACK_LIST
     network.to_float(mstype.float16)
     cells = network.name_cells()
     change = False
@@ -99,7 +257,7 @@ def _auto_black_list(network, black_list=None):
         subcell = cells[name]
         if subcell == network:
             continue
-        if isinstance(subcell, black_list):
+        if isinstance(subcell, tuple(black_list)):
             network._cells[name] = _OutputTo16(subcell.to_float(mstype.float32))
             change = True
         else:
@@ -117,7 +275,7 @@ def auto_mixed_precision(network, amp_level="O0"):
         amp_level (str): Supports ["O0", "O1", "O2", "O3"]. Default: "O0".
 
             - "O0": Do not change.
-            - "O1": [DEMO] Cast the operators in white_list to float16, the remaining operators are kept in float32.
+            - "O1": Cast the operators in white_list to float16, the remaining operators are kept in float32.
             - "O2": Cast network to float16, keep operators in black_list run in float32,
             - "O3": Cast network to float16.
 
@@ -125,25 +283,24 @@ def auto_mixed_precision(network, amp_level="O0"):
         ValueError: If amp level is not supported.
 
     Examples:
-        >>> from mindpsore import amp, nn
+        >>> from mindspore import amp, nn
         >>> network = LeNet5()
         >>> amp_level = "O1"
         >>> net = amp.auto_mixed_precision(network, amp_level)
     """
     if not isinstance(network, nn.Cell):
         raise TypeError("The network type should be Cell.")
+
     if amp_level == "O0":
         pass
     elif amp_level == "O1":
-        _auto_white_list(network)
-        return network
+        return _auto_white_list(network, AMP_WHITE_LIST)
     elif amp_level == "O2":
-        _auto_black_list(network)
+        _auto_black_list(network, AMP_BLACK_LIST)
     elif amp_level == "O3":
         network.to_float(mstype.float16)
     else:
         raise ValueError("The amp level {} is not supported".format(amp_level))
-
     if amp_level in ("O2", "O3"):
         network = _OutputTo32(network)
     return network
@@ -157,7 +314,7 @@ def _do_keep_batchnorm_fp32(network):
         subcell = cells[name]
         if subcell == network:
             continue
-        elif isinstance(subcell, AMP_BLACK_LIST):
+        elif isinstance(subcell, nn.Cell) and isinstance(subcell, tuple(AMP_BLACK_LIST)):
             network._cells[name] = _OutputTo16(subcell.to_float(mstype.float32))
             change = True
         else:
@@ -208,8 +365,8 @@ def _check_level(level, boost_level):
     if not isinstance(level, str):
         raise TypeError("The argument `level` must be a string in ['O0', 'O1', 'O2', 'O3', 'auto'], \
                          but got type {}.".format(type(level)))
-    validator.check('level', level, "", ['O0', 'O1', 'O2', 'O3', 'auto'], Rel.IN)
-    validator.check('boost_level', boost_level, "", ['O0', 'O1', 'O2'], Rel.IN)
+    validator.check('level', level, "", ['O0', 'O1', 'O2', 'O3', 'auto'], validator.IN)
+    validator.check('boost_level', boost_level, "", ['O0', 'O1', 'O2'], validator.IN)
 
     if level == "auto":
         device_target = context.get_context('device_target')
@@ -231,13 +388,12 @@ def _add_loss_network(network, loss_fn, cast_model_type):
     """Add loss network."""
 
     class WithLossCell(nn.Cell):
-        "Wrap loss for amp. Cast network output back to float32"
-
+        """Wrap loss for amp. Cast network output back to float32."""
         def __init__(self, backbone, loss_fn):
             super(WithLossCell, self).__init__(auto_prefix=False)
             self._backbone = backbone
             self._loss_fn = loss_fn
-            if backbone.jit_config_dict:
+            if isinstance(backbone, nn.Cell) and backbone.jit_config_dict:
                 self._jit_config_dict = backbone.jit_config_dict
 
         def construct(self, data, label):
@@ -265,7 +421,9 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
         level (str): Supports ["O0", "O1", "O2", "O3", "auto"]. Default: "O0".
 
             - "O0": Do not change.
-            - "O1": [DEMO] Cast the operators in white_list to float16, the remaining operators are kept in float32.
+            - "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.
             - "O3": Cast network to float16, with additional property `keep_batchnorm_fp32=False` .
@@ -302,7 +460,7 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
             (with property `drop_overflow_update=False` ).
 
     Examples:
-        >>> from mindpsore import amp, nn
+        >>> from mindspore import amp, nn
         >>> network = LeNet5()
         >>> net_loss = nn.SoftmaxCrossEntropyWithLogits(reduction="mean")
         >>> net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)
@@ -327,7 +485,7 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
     elif config["cast_model_type"] == mstype.float32 and level in ("O2", "O3"):
         pass
     else:
-        auto_mixed_precision(network, level)
+        network = auto_mixed_precision(network, level)
 
     if loss_fn:
         network = _add_loss_network(network, loss_fn, config["cast_model_type"])
@@ -360,3 +518,116 @@ def build_train_network(network, optimizer, loss_fn=None, level='O0', boost_leve
     else:
         network = nn.TrainOneStepCell(network, optimizer, loss_scale).set_train()
     return network
+
+
+def get_white_list():
+    """
+    Provide a copy of internal white list used by auto mixed precision.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Returns:
+        list, A copy of internal white list.
+    """
+    white_list = AMP_WHITE_LIST.copy()
+    return white_list
+
+
+def get_black_list():
+    """
+    Provide a copy of internal black list used by auto mixed precision.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Returns:
+        list, A copy of internal black list.
+    """
+    black_list = AMP_BLACK_LIST.copy()
+    return black_list
+
+
+def custom_mixed_precision(network, *, white_list=None, black_list=None):
+    """
+    Custom mixed precision by setting whitelist or blacklist.
+    When the `white_list` is provided, primitives and cells in `white_list` will perform the precision conversion.
+    When the `black_list` is provided, primitives and cells that are not in `black_list` will perform the pereision
+    conversion.
+    Only one of `white_list` and `black_list` should be provided.
+
+    .. warning::
+        This is an experimental API that is subject to change or deletion.
+
+    Note:
+        - `custom_mixed_precision` should not be used at the same time as `auto_mixed_precision` . When both
+          `build_train_network` and `custom_mixed_precision` are used, `build_train_network` need to be called with
+          `level='O0'` before call `custom_mixed_precision` .
+        - Primitives for blacklist is not support yet.
+
+    Args:
+        network (Cell): Definition of the network.
+        white_list (list[Primitive, Cell], optional): White list of custom mixed precision. Defaults: None, means
+            white list is not used.
+        black_list (list[Primitive, Cell], optional): Black list of custom mixed precision. Defaults: None, means
+            black list is not used.
+
+    Returns:
+        network (Cell), A network supporting mixed precision.
+
+    Raises:
+        TypeError: The network type is not Cell.
+        ValueError: Neither `white_list` nor `black_list` is provided.
+        ValueError: Both `white_list` and `black_list` are provided.
+
+    Examples:
+        >>> from mindspore import amp
+        >>> net = MyNet()
+        >>> custom_white_list = amp.get_white_list()
+        >>> custom_white_list.append(nn.Tanhshrink)
+        >>> net = amp.custom_mixed_precision(net, white_list=custom_white_list)
+    """
+    if not isinstance(network, nn.Cell):
+        raise TypeError("The network type should be Cell.")
+
+    if white_list is None and black_list is None:
+        raise ValueError("For custom_mixed_precision, one of white_list and black_list must be provided.")
+
+    if white_list is not None and black_list is not None:
+        raise ValueError("For custom_mixed_precision, the white_list or black_list cannot be provided "
+                         "at the same time, please provide one or the other.")
+
+    if white_list is not None:
+        _list_check(white_list, "white_list")
+        return _auto_white_list(network, white_list)
+
+    _list_check(black_list, "black_list")
+    _auto_black_list(network, black_list)
+    network = _OutputTo32(network)
+    return network
+
+
+def _list_check(custom_list: list, list_name: str):
+    """
+    check whether custom list is valid
+
+    Raises:
+        TypeError: The type of custom_list is not list.
+        TypeError: The element in custom_list is not a class.
+        TypeError: The subclass of element in custom_list is not one of ['Cell', 'Primitive'].
+    """
+    if not isinstance(custom_list, list):
+        raise TypeError(f"The type of {list_name} should be list, but got {type(custom_list)}")
+
+    for elem in custom_list:
+        if not isinstance(elem, type):
+            raise TypeError(f"The element in {list_name} should be a class, but got {elem}")
+
+        if not issubclass(elem, nn.Cell) and not issubclass(elem, Primitive):
+            raise TypeError(f"The subclass of element in {list_name} should be one of 'Cell' and 'Primitive', "
+                            f"but got {elem}")
+
+    if list_name == 'black_list':
+        for elem in AMP_BLACK_LIST:
+            if elem not in custom_list:
+                logger.warning(f"{elem} is removed from internal black list.")

mindspore/train/callback/__init__.py

@@ -33,7 +33,9 @@ from mindspore.train.callback._history import History
 from mindspore.train.callback._lambda_callback import LambdaCallback
 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
 
 __all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint",
            "SummaryCollector", "CheckpointConfig", "RunContext", "LearningRateScheduler", "SummaryLandscape",
-           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping"]
+           "History", "LambdaCallback", "ReduceLROnPlateau", "EarlyStopping", "OnRequestExit", "BackupAndRestore"]