Add a callback named SummaryCollector and delete SummaryStep callback

I added a SummaryCollector to help users automatically collect information
such as the network, loss, learning rate and so on, making it easier to collect this information.

It also can collect train lineage and eval lineage information which is
collected by TrainLineage Callback and EvalLineage Callback in
MindInsight.

I also add some UT for SummaryCollect to keep the code correct.

mindspore/train/_utils.py

@@ -14,7 +14,10 @@
 # ============================================================================
 """Train utility."""
 import os
+from collections.abc import Iterable
+
 import numpy as np
+
 from mindspore.common.tensor import Tensor
 from mindspore.common.dtype import dtype_to_nptype, pytype_to_dtype
 from mindspore.common import dtype as mstype
@@ -213,6 +216,7 @@ def _check_to_numpy(plugin, tensor):
         raise ValueError('The tensor should not be empty.')
     return np_value
 
+
 def _check_lineage_value(plugin, value):
     """Check the lineage value."""
     def raises(plugin, prototype):
@@ -229,3 +233,20 @@ def _check_lineage_value(plugin, value):
 
     if plugin == 'custom_lineage_data' and not isinstance(value, UserDefinedInfo):
         raises(plugin, UserDefinedInfo)
+
+
+def check_value_type(arg_name, arg_value, valid_types):
+    """Checks whether a value is instance of some types."""
+    valid_types = tuple(valid_types) if isinstance(valid_types, Iterable) else (valid_types,)
+    is_valid = True
+
+    # bool is subclass of int, so for a bool value, we need to extra check
+    if isinstance(arg_value, int) and isinstance(arg_value, bool) and bool not in valid_types:
+        is_valid = False
+
+    if not isinstance(arg_value, valid_types):
+        is_valid = False
+
+    if not is_valid:
+        raise TypeError(f'For `{arg_name}` the type should be a valid type of {[t.__name__ for t in valid_types]}, '
+                        f'bug got {type(arg_value).__name__}.')

mindspore/train/callback/__init__.py

@@ -22,7 +22,8 @@ from ._checkpoint import CheckpointConfig
 from ._checkpoint import CheckpointManager as _CheckpointManager
 from ._checkpoint import ModelCheckpoint
 from ._loss_monitor import LossMonitor
-from ._summary_step import SummaryStep
 from ._time_monitor import TimeMonitor
+from ._summary_collector import SummaryCollector
 
-__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint", "SummaryStep", "CheckpointConfig", "RunContext"]
+__all__ = ["Callback", "LossMonitor", "TimeMonitor", "ModelCheckpoint",
+           "SummaryCollector", "CheckpointConfig", "RunContext"]

mindspore/train/callback/_dataset_graph.py

+# Copyright 2020 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.
+# ============================================================================
+"""Define dataset graph related operations."""
+import json
+from importlib import import_module
+
+from mindspore.train import lineage_pb2
+
+
+class DatasetGraph:
+    """Handle the data graph and packages it into binary data."""
+    def package_dataset_graph(self, dataset):
+        """
+        packages dataset graph into binary data
+
+        Args:
+            dataset (MindData): refer to MindDataset
+
+        Returns:
+            DatasetGraph, a object of lineage_pb2.DatasetGraph.
+        """
+        dataset_package = import_module('mindspore.dataset')
+        dataset_dict = dataset_package.serialize(dataset)
+        json_str = json.dumps(dataset_dict, indent=2)
+        dataset_dict = json.loads(json_str)
+        dataset_graph_proto = lineage_pb2.DatasetGraph()
+        if "children" in dataset_dict:
+            children = dataset_dict.pop("children")
+            if children:
+                self._package_children(children=children, message=dataset_graph_proto)
+            self._package_current_dataset(operation=dataset_dict, message=dataset_graph_proto)
+        return dataset_graph_proto
+
+    def _package_children(self, children, message):
+        """
+        Package children in dataset operation.
+
+        Args:
+            children (list[dict]): Child operations.
+            message (DatasetGraph): Children proto message.
+        """
+        for child in children:
+            if child:
+                child_graph_message = getattr(message, "children").add()
+                grandson = child.pop("children")
+                if grandson:
+                    self._package_children(children=grandson, message=child_graph_message)
+                # package other parameters
+                self._package_current_dataset(operation=child, message=child_graph_message)
+
+    def _package_current_dataset(self, operation, message):
+        """
+        Package operation parameters in event message.
+
+        Args:
+            operation (dict): Operation dict.
+            message (Operation): Operation proto message.
+        """
+        for key, value in operation.items():
+            if value and key == "operations":
+                for operator in value:
+                    self._package_enhancement_operation(
+                        operator,
+                        message.operations.add()
+                    )
+            elif value and key == "sampler":
+                self._package_enhancement_operation(
+                    value,
+                    message.sampler
+                )
+            else:
+                self._package_parameter(key, value, message.parameter)
+
+    def _package_enhancement_operation(self, operation, message):
+        """
+        Package enhancement operation in MapDataset.
+
+        Args:
+            operation (dict): Enhancement operation.
+            message (Operation): Enhancement operation proto message.
+        """
+        for key, value in operation.items():
+            if isinstance(value, list):
+                if all(isinstance(ele, int) for ele in value):
+                    message.size.extend(value)
+                else:
+                    message.weights.extend(value)
+            else:
+                self._package_parameter(key, value, message.operationParam)
+
+    @staticmethod
+    def _package_parameter(key, value, message):
+        """
+        Package parameters in operation.
+
+        Args:
+            key (str): Operation name.
+            value (Union[str, bool, int, float, list, None]): Operation args.
+            message (OperationParameter): Operation proto message.
+        """
+        if isinstance(value, str):
+            message.mapStr[key] = value
+        elif isinstance(value, bool):
+            message.mapBool[key] = value
+        elif isinstance(value, int):
+            message.mapInt[key] = value
+        elif isinstance(value, float):
+            message.mapDouble[key] = value
+        elif isinstance(value, list) and key != "operations":
+            if value:
+                replace_value_list = list(map(lambda x: "" if x is None else x, value))
+                message.mapStrList[key].strValue.extend(replace_value_list)
+        elif value is None:
+            message.mapStr[key] = "None"
+        else:
+            raise ValueError(f"Parameter {key} is not supported in event package.")

mindspore/train/model.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ============================================================================
 """Model."""
+from collections.abc import Iterable
+
 import numpy as np
 
 from mindspore import log as logger
@@ -345,7 +347,8 @@ class Model:
         cb_params.parallel_mode = self._parallel_mode
         cb_params.device_number = self._device_number
         cb_params.train_dataset = train_dataset
-        cb_params.list_callback = callbacks
+        cb_params.list_callback = self._transform_callbacks(callbacks)
+        cb_params.train_dataset_element = None
 
         # build callback list
         with _CallbackManager(callbacks) as list_callback:
@@ -358,6 +361,17 @@ class Model:
             else:
                 self._train_dataset_sink_process(epoch, train_dataset, list_callback, cb_params)
 
+    @staticmethod
+    def _transform_callbacks(callbacks):
+        """Transform callback to a list."""
+        if callbacks is None:
+            return []
+
+        if isinstance(callbacks, Iterable):
+            return list(callbacks)
+
+        return [callbacks]
+
     def _train_dataset_sink_process(self, epoch, train_dataset, list_callback=None, cb_params=None):
         """
         Training process. The data would be passed to network through dataset channel.
@@ -449,6 +463,7 @@ class Model:
                     scaling_sens = self._get_scaling_sens()
                     next_element = tuple(next_element) + (Tensor(scaling_sens, mstype.float32),)
 
+                cb_params.train_dataset_element = next_element
                 outputs = self._train_network(*next_element)
                 cb_params.net_outputs = outputs
                 if self._loss_scale_manager and self._loss_scale_manager.get_drop_overflow_update():
@@ -628,6 +643,7 @@ class Model:
         cb_params.batch_num = valid_dataset.get_dataset_size()
         cb_params.mode = "eval"
         cb_params.cur_step_num = 0
+        cb_params.list_callback = self._transform_callbacks(callbacks)
 
         self._eval_network.set_train(mode=False)
         self._eval_network.phase = 'eval'

mindspore/train/callback/_summary_step.py → mindspore/train/summary/enum.py

@@ -12,45 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ============================================================================
-"""SummaryStep Callback class."""
+"""Summary's enumeration file."""
+from enum import Enum
 
-from ._callback import Callback
 
+class BaseEnum(Enum):
+    """The base enum class."""
 
-class SummaryStep(Callback):
-    """
-    The summary callback class.
+    @classmethod
+    def to_list(cls):
+        """Converts the enumeration into a list."""
+        return [member.value for member in cls.__members__.values()]
 
-    Args:
-        summary (Object): Summary recode object.
-        flush_step (int): Number of interval steps to execute. Default: 10.
-    """
 
-    def __init__(self, summary, flush_step=10):
-        super(SummaryStep, self).__init__()
-        if not isinstance(flush_step, int) or isinstance(flush_step, bool) or flush_step <= 0:
-            raise ValueError("`flush_step` should be int and greater than 0")
-        self._summary = summary
-        self._flush_step = flush_step
+class PluginEnum(BaseEnum):
+    """The list of plugins currently supported by the summary."""
+    GRAPH = 'graph'
+    SCALAR = 'scalar'
+    IMAGE = 'image'
+    TENSOR = 'tensor'
+    HISTOGRAM = 'histogram'
+    TRAIN_LINEAGE = 'train_lineage'
+    EVAL_LINEAGE = 'eval_lineage'
+    DATASET_GRAPH = 'dataset_graph'
 
-    def __enter__(self):
-        self._summary.__enter__()
-        return self
 
-    def __exit__(self, *err):
-        return self._summary.__exit__(*err)
-
-    def step_end(self, run_context):
-        """
-        Save summary.
-
-        Args:
-            run_context (RunContext): Context of the train running.
-        """
-        cb_params = run_context.original_args()
-        if cb_params.cur_step_num % self._flush_step == 0:
-            self._summary.record(cb_params.cur_step_num, cb_params.train_network)
-
-    @property
-    def summary_file_name(self):
-        return self._summary.full_file_name
+class ModeEnum(BaseEnum):
+    """The modes currently supported by the summary."""
+    TRAIN = 'train'
+    EVAL = 'eval'

tests/st/summary/test_gpu_summary.py

@@ -75,7 +75,7 @@ class TestGpuSummary:
         if not os.path.exists(self.summary_dir):
             os.mkdir(self.summary_dir)
 
-    def teardown_emthod(self):
+    def teardown_method(self):
         """Run after method."""
         if os.path.exists(self.summary_dir):
             shutil.rmtree(self.summary_dir)

tests/ut/python/train/summary/test_graph_summary.py

@@ -20,8 +20,8 @@ import numpy as np
 import mindspore.nn as nn
 from mindspore import Model, context
 from mindspore.nn.optim import Momentum
-from mindspore.train.callback import SummaryStep
-from mindspore.train.summary.summary_record import SummaryRecord
+from mindspore.train.summary import SummaryRecord
+from mindspore.train.callback import SummaryCollector
 from .....dataset_mock import MindData
 
 CUR_DIR = os.getcwd()
@@ -107,16 +107,9 @@ def test_graph_summary_sample():
     model = Model(net, loss_fn=loss, optimizer=optim, metrics=None)
     with SummaryRecord(SUMMARY_DIR, file_suffix="_MS_GRAPH", network=model._train_network) as test_writer:
         model.train(2, dataset)
-        # step 2: create the Event
         for i in range(1, 5):
             test_writer.record(i)
 
-        # step 3: send the event to mq
-
-        # step 4: accept the event and write the file
-
-        log.debug("finished test_graph_summary_sample")
-
 
 def test_graph_summary_callback():
     dataset = get_dataset()
@@ -125,18 +118,8 @@ def test_graph_summary_callback():
     optim = Momentum(net.trainable_params(), 0.1, 0.9)
     context.set_context(mode=context.GRAPH_MODE)
     model = Model(net, loss_fn=loss, optimizer=optim, metrics=None)
-    with SummaryRecord(SUMMARY_DIR, file_suffix="_MS_GRAPH", network=model._train_network) as test_writer:
-        summary_cb = SummaryStep(test_writer, 1)
-        model.train(2, dataset, callbacks=summary_cb)
-
-
-def test_graph_summary_callback2():
-    dataset = get_dataset()
-    net = Net()
-    loss = nn.SoftmaxCrossEntropyWithLogits()
-    optim = Momentum(net.trainable_params(), 0.1, 0.9)
-    context.set_context(mode=context.GRAPH_MODE)
-    model = Model(net, loss_fn=loss, optimizer=optim, metrics=None)
-    with SummaryRecord(SUMMARY_DIR, file_suffix="_MS_GRAPH", network=net) as test_writer:
-        summary_cb = SummaryStep(test_writer, 1)
-        model.train(2, dataset, callbacks=summary_cb)
+    summary_collector = SummaryCollector(SUMMARY_DIR,
+                                         collect_freq=1,
+                                         keep_default_action=False,
+                                         collect_specified_data={'collect_graph': True})
+    model.train(1, dataset, callbacks=[summary_collector])

tests/ut/python/train/summary/test_image_summary.py

@@ -26,9 +26,8 @@ import mindspore.nn as nn
 from mindspore import Model, context
 from mindspore import Tensor
 from mindspore.nn.optim import Momentum
-from mindspore.train.callback import SummaryStep
-from mindspore.train.summary.summary_record import SummaryRecord, \
-    _cache_summary_tensor_data
+from mindspore.train.summary.summary_record import SummaryRecord, _cache_summary_tensor_data
+from mindspore.train.callback import Callback
 from .....dataset_mock import MindData
 
 CUR_DIR = os.getcwd()
@@ -155,7 +154,8 @@ def get_dataset():
     return dataset
 
 
-class ImageSummaryCallback:
+class ImageSummaryCallback(Callback):
+    """Image summary callback."""
 
     def __init__(self, summary_record):
         self._summary_record = summary_record
@@ -164,9 +164,10 @@ class ImageSummaryCallback:
         return self
 
     def __exit__(self, *err):
-        pass
+        self._summary_record.close()
 
     def record(self, step, train_network=None):
+        """record data."""
         self._summary_record.record(step, train_network)
         self._summary_record.flush()
 
@@ -183,9 +184,8 @@ def test_image_summary_train():
         # step 2: create the Event
 
         model = get_model()
-        fn = ImageSummaryCallback(test_writer)
-        summary_recode = SummaryStep(fn, 1)
-        model.train(2, dataset, callbacks=summary_recode)
+        callback = ImageSummaryCallback(test_writer)
+        model.train(2, dataset, callbacks=[callback])
 
         # step 3: send the event to mq
 

tests/ut/python/train/summary/test_summary.py

@@ -24,11 +24,9 @@ import random
 import numpy as np
 import pytest
 
-
 import mindspore.nn as nn
 from mindspore.common.tensor import Tensor
 from mindspore.ops import operations as P
-from mindspore.train.callback import SummaryStep
 from mindspore.train.summary.summary_record import SummaryRecord, _cache_summary_tensor_data
 
 CUR_DIR = os.getcwd()
@@ -192,16 +190,6 @@ def test_scalar_summary_with_ge_2():
 
 def test_validate():
     with SummaryRecord(SUMMARY_DIR) as sr:
-        with pytest.raises(ValueError):
-            SummaryStep(sr, 0)
-        with pytest.raises(ValueError):
-            SummaryStep(sr, -1)
-        with pytest.raises(ValueError):
-            SummaryStep(sr, 1.2)
-        with pytest.raises(ValueError):
-            SummaryStep(sr, True)
-        with pytest.raises(ValueError):
-            SummaryStep(sr, "str")
         sr.record(1)
         with pytest.raises(ValueError):
             sr.record(False)
@@ -215,17 +203,3 @@ def test_validate():
             sr.record("str")
         with pytest.raises(ValueError):
             sr.record(sr)
-
-    SummaryStep(sr, 1)
-    with pytest.raises(ValueError):
-        SummaryStep(sr, 1.2)
-    with pytest.raises(ValueError):
-        SummaryStep(sr, False)
-    with pytest.raises(ValueError):
-        SummaryStep(sr, "str")
-    with pytest.raises(ValueError):
-        SummaryStep(sr, (1, 2))
-    with pytest.raises(ValueError):
-        SummaryStep(sr, [3, 4])
-    with pytest.raises(ValueError):
-        SummaryStep(sr, sr)

tests/ut/python/train/summary/test_summary_abnormal_input.py

@@ -59,7 +59,8 @@ def test_summaryrecord_input_null_string():
     log.debug("begin test_summaryrecord_input_null_string")
     # step 0: create the thread
     try:
-        SummaryRecord("")
+        with SummaryRecord(""):
+            pass
     except:
         assert True
     else:
@@ -71,7 +72,8 @@ def test_summaryrecord_input_None():
     log.debug("begin test_summaryrecord_input_None")
     # step 0: create the thread
     try:
-        SummaryRecord(None)
+        with SummaryRecord(None):
+            pass
     except:
         assert True
     else:
@@ -83,7 +85,8 @@ def test_summaryrecord_input_relative_dir_1():
     log.debug("begin test_summaryrecord_input_relative_dir_1")
     # step 0: create the thread
     try:
-        SummaryRecord("./test_temp_summary_event_file/")
+        with SummaryRecord("./test_temp_summary_event_file/"):
+            pass
     except:
         assert False
     else:
@@ -95,7 +98,8 @@ def test_summaryrecord_input_relative_dir_2():
     log.debug("begin test_summaryrecord_input_relative_dir_2")
     # step 0: create the thread
     try:
-        SummaryRecord("../summary/")
+        with SummaryRecord("../summary/"):
+            pass
     except:
         assert False
     else:
@@ -107,7 +111,8 @@ def test_summaryrecord_input_invalid_type_dir():
     log.debug("begin test_summaryrecord_input_invalid_type_dir")
     # step 0: create the thread
     try:
-        SummaryRecord(32)
+        with SummaryRecord(32):
+            pass
     except:
         assert True
     else:
@@ -119,7 +124,8 @@ def test_mulit_layer_directory():
     log.debug("begin test_mulit_layer_directory")
     # step 0: create the thread
     try:
-        SummaryRecord("./test_temp_summary_event_file/test/t1/")
+        with SummaryRecord("./test_temp_summary_event_file/test/t1/"):
+            pass
     except:
         assert False
     else:

tests/ut/python/train/summary/test_summary_collector.py

+# Copyright 2020 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.
+# ============================================================================
+"""Test the exception parameter scenario for summary collector."""
+import os
+import tempfile
+import shutil
+import pytest
+
+from mindspore.train.callback import SummaryCollector
+
+
+class TestSummaryCollector:
+    """Test the exception parameter for summary collector."""
+    base_summary_dir = ''
+
+    def setup_class(self):
+        """Run before test this class."""
+        self.base_summary_dir = tempfile.mkdtemp(suffix='summary')
+
+    def teardown_class(self):
+        """Run after test this class."""
+        if os.path.exists(self.base_summary_dir):
+            shutil.rmtree(self.base_summary_dir)
+
+    @pytest.mark.parametrize("summary_dir", [1234, None, True, ''])
+    def test_params_with_summary_dir_value_error(self, summary_dir):
+        """Test the exception scenario for summary dir."""
+        if isinstance(summary_dir, str):
+            with pytest.raises(ValueError) as exc:
+                SummaryCollector(summary_dir=summary_dir)
+            assert str(exc.value) == 'For `summary_dir` the value should be a valid string of path, ' \
+                                     'but got empty string.'
+        else:
+            with pytest.raises(TypeError) as exc:
+                SummaryCollector(summary_dir=summary_dir)
+            assert 'For `summary_dir` the type should be a valid type' in str(exc.value)
+
+    def test_params_with_summary_dir_not_dir(self):
+        """Test the given summary dir parameter is not a directory."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        summary_file = os.path.join(summary_dir, 'temp_file.txt')
+        with open(summary_file, 'w') as file_handle:
+            file_handle.write('temp')
+        print(os.path.isfile(summary_file))
+        with pytest.raises(NotADirectoryError):
+            SummaryCollector(summary_dir=summary_file)
+
+    @pytest.mark.parametrize("collect_freq", [None, 0, 0.01])
+    def test_params_with_collect_freq_exception(self, collect_freq):
+        """Test the exception scenario for collect freq."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        if isinstance(collect_freq, int):
+            with pytest.raises(ValueError) as exc:
+                SummaryCollector(summary_dir=summary_dir, collect_freq=collect_freq)
+            expected_msg = f'For `collect_freq` the value should be greater than 0, but got `{collect_freq}`.'
+            assert expected_msg == str(exc.value)
+        else:
+            with pytest.raises(TypeError) as exc:
+                SummaryCollector(summary_dir=summary_dir, collect_freq=collect_freq)
+            expected_msg = f"For `collect_freq` the type should be a valid type of ['int'], " \
+                           f'bug got {type(collect_freq).__name__}.'
+            assert expected_msg == str(exc.value)
+
+    @pytest.mark.parametrize("action", [None, 123, '', '123'])
+    def test_params_with_action_exception(self, action):
+        """Test the exception scenario for action."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        with pytest.raises(TypeError) as exc:
+            SummaryCollector(summary_dir=summary_dir, keep_default_action=action)
+        expected_msg = f"For `keep_default_action` the type should be a valid type of ['bool'], " \
+                       f"bug got {type(action).__name__}."
+        assert expected_msg == str(exc.value)
+
+    @pytest.mark.parametrize("collect_specified_data", [123])
+    def test_params_with_collect_specified_data_type_error(self, collect_specified_data):
+        """Test type error scenario for collect specified data param."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        with pytest.raises(TypeError) as exc:
+            SummaryCollector(summary_dir, collect_specified_data=collect_specified_data)
+
+        expected_msg = f"For `collect_specified_data` the type should be a valid type of ['dict', 'NoneType'], " \
+                       f"bug got {type(collect_specified_data).__name__}."
+
+        assert expected_msg == str(exc.value)
+
+    @pytest.mark.parametrize("collect_specified_data", [
+        {
+            123: 123
+        },
+        {
+            None: True
+        }
+    ])
+    def test_params_with_collect_specified_data_key_type_error(self, collect_specified_data):
+        """Test the key of collect specified data param."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        with pytest.raises(TypeError) as exc:
+            SummaryCollector(summary_dir, collect_specified_data=collect_specified_data)
+
+        param_name = list(collect_specified_data)[0]
+        expected_msg = f"For `{param_name}` the type should be a valid type of ['str'], " \
+                       f"bug got {type(param_name).__name__}."
+        assert expected_msg == str(exc.value)
+
+    @pytest.mark.parametrize("collect_specified_data", [
+        {
+            'collect_metric': None
+        },
+        {
+            'collect_graph': 123
+        },
+        {
+            'histogram_regular': 123
+        },
+    ])
+    def test_params_with_collect_specified_data_value_type_error(self, collect_specified_data):
+        """Test the value of collect specified data param."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        with pytest.raises(TypeError) as exc:
+            SummaryCollector(summary_dir, collect_specified_data=collect_specified_data)
+
+        param_name = list(collect_specified_data)[0]
+        param_value = collect_specified_data[param_name]
+        expected_type = "['bool']" if param_name != 'histogram_regular' else "['str', 'NoneType']"
+        expected_msg = f'For `{param_name}` the type should be a valid type of {expected_type}, ' \
+                       f'bug got {type(param_value).__name__}.'
+
+        assert expected_msg == str(exc.value)
+
+    def test_params_with_collect_specified_data_unexpected_key(self):
+        """Test the collect_specified_data parameter with unexpected key."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        data = {'unexpected_key': True}
+        with pytest.raises(ValueError) as exc:
+            SummaryCollector(summary_dir, collect_specified_data=data)
+        expected_msg = f"For `collect_specified_data` the keys {set(data)} are unsupported."
+        assert expected_msg == str(exc.value)
+
+    @pytest.mark.parametrize("custom_lineage_data", [
+        123,
+        {
+            'custom': {}
+        },
+        {
+            'custom': None
+        },
+        {
+            123: 'custom'
+        }
+    ])
+    def test_params_with_custom_lineage_data_type_error(self, custom_lineage_data):
+        """Test the custom lineage data parameter type error."""
+        summary_dir = tempfile.mkdtemp(dir=self.base_summary_dir)
+        with pytest.raises(TypeError) as exc:
+            SummaryCollector(summary_dir, custom_lineage_data=custom_lineage_data)
+
+        if not isinstance(custom_lineage_data, dict):
+            expected_msg = f"For `custom_lineage_data` the type should be a valid type of ['dict', 'NoneType'], " \
+                           f"bug got {type(custom_lineage_data).__name__}."
+        else:
+            param_name = list(custom_lineage_data)[0]
+            param_value = custom_lineage_data[param_name]
+            if not isinstance(param_name, str):
+                arg_name = f'custom_lineage_data -> {param_name}'
+                expected_msg = f"For `{arg_name}` the type should be a valid type of ['str'], " \
+                               f'bug got {type(param_name).__name__}.'
+            else:
+                arg_name = f'the value of custom_lineage_data -> {param_name}'
+                expected_msg = f"For `{arg_name}` the type should be a valid type of ['int', 'str', 'float'], " \
+                               f'bug got {type(param_value).__name__}.'
+
+        assert expected_msg == str(exc.value)

tests/ut/python/train/test_training.py

@@ -20,8 +20,8 @@ import pytest
 import mindspore.nn as nn
 from mindspore import Model, context
 from mindspore import Tensor
+from mindspore.train.callback import Callback
 from mindspore.nn.optim import Momentum
-from mindspore.train.callback import SummaryStep
 from ..ut_filter import non_graph_engine
 from ....dataset_mock import MindData
 
@@ -174,7 +174,7 @@ class TestGraphMode:
         model.train(1, dataset)
 
 
-class CallbackTest:
+class CallbackTest(Callback):
     """ CallbackTest definition """
 
     def __init__(self):
@@ -186,19 +186,19 @@ class CallbackTest:
     def __exit__(self, *err):
         pass
 
-    def record(self, step, *args):
-        print(step, args)
+    def step_end(self, run_context):
+        cb_params = run_context.original_args()
+        print(cb_params.cur_epoch_num, cb_params.cur_step_num)
 
 
 def test_train_callback(test_with_simu):
     """ test_train_callback """
     dataset = get_dataset()
     model = get_model()
-    fn = CallbackTest()
-    summary_recode = SummaryStep(fn, 2)
+    callback = CallbackTest()
     if test_with_simu:
         return
-    model.train(2, dataset, callbacks=summary_recode)
+    model.train(2, dataset, callbacks=callback)
 
 
 log = logging.getLogger("test")