# Copyright 2019 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. # ============================================================================ """Unittest for model_lineage.py""" import os import shutil import unittest from unittest import mock, TestCase from unittest.mock import MagicMock from mindinsight.lineagemgr.collection.model.model_lineage import TrainLineage, EvalLineage, \ AnalyzeObject from mindinsight.lineagemgr.common.exceptions.exceptions import \ LineageLogError, LineageGetModelFileError, MindInsightException from mindspore.common.tensor import Tensor from mindspore.dataset.engine import MindDataset, Dataset from mindspore.nn import Optimizer, WithLossCell, TrainOneStepWithLossScaleCell, \ SoftmaxCrossEntropyWithLogits from mindspore.train.callback import RunContext, ModelCheckpoint, SummaryStep from mindspore.train.summary import SummaryRecord @mock.patch('builtins.open') @mock.patch('os.makedirs') class TestModelLineage(TestCase): """Test TrainLineage and EvalLineage class in model_lineage.py.""" @classmethod def setUpClass(cls): cls.lineage_list = ['train_network', 'loss_fn', 'optimizer', 'train_dataset', 'valid_dataset', 'epoch', 'valid_step', 'hybrid_parallel', 'data_parallel_size', 'auto_parallel', 'device_number', 'batch_num', 'summary_log_path', 'model_ckpt'] cls.run_context = {key: None for key in cls.lineage_list} cls.run_context['net_outputs'] = Tensor() cls.my_run_context = RunContext cls.my_train_module = TrainLineage cls.my_eval_module = EvalLineage cls.my_analyze_module = AnalyzeObject cls.my_summary_record = SummaryRecord cls.summary_log_path = '/path/to/summary_log' @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_summary_record_exception(self, *args): """Test SummaryRecord with exception.""" args[0].return_value = None summary_record = self.my_summary_record(self.summary_log_path) with self.assertRaises(MindInsightException) as context: self.my_train_module(summary_record=summary_record, raise_exception=1) self.assertTrue(f'Invalid value for raise_exception.' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.ds') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.LineageSummary.record_dataset_graph') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_optimizer_by_network') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_optimizer') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_network') def test_begin(self, *args): """Test TrainLineage.begin method.""" args[1].return_value = None args[2].return_value = Optimizer(Tensor(0.1)) args[3].return_value = None args[5].serialize.return_value = {} run_context = {'optimizer': Optimizer(Tensor(0.1)), 'epoch_num': 10} train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path)) train_lineage.begin(self.my_run_context(run_context)) args[4].assert_called() @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.ds') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.LineageSummary.record_dataset_graph') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_optimizer_by_network') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_optimizer') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_network') def test_begin_error(self, *args): """Test TrainLineage.begin method.""" args[1].return_value = None args[2].return_value = Optimizer(Tensor(0.1)) args[3].return_value = None args[4].side_effect = Exception args[5].serialize.return_value = {} run_context = {'optimizer': Optimizer(Tensor(0.1)), 'epoch_num': 10} train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaisesRegex(LineageLogError, 'Dataset graph log error'): train_lineage.begin(self.my_run_context(run_context)) train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path)) train_lineage.begin(self.my_run_context(run_context)) args[4].assert_called() @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_begin_exception(self, *args): """Test TrainLineage.begin method with exception.""" args[0].return_value = None train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaises(Exception) as context: train_lineage.begin(self.run_context) self.assertTrue('Invalid TrainLineage run_context.' in str(context.exception)) run_context = {key: None for key in self.lineage_list} run_context['optimizer'] = 1 with self.assertRaises(Exception) as context: train_lineage.begin(self.my_run_context(run_context)) self.assertTrue('The parameter optimizer is invalid.' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_model_size') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_file_path') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.LineageSummary.record_train_lineage') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_dataset') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_optimizer') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_network') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_train_run_context') @mock.patch('builtins.float') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_train_end(self, *args): """Test TrainLineage.end method.""" args[1].return_value = 2.0 args[2].return_value = True args[3].return_value = True args[4].return_value = None args[5].return_value = None args[6].return_value = None args[7].return_value = (None, None) args[8].return_value = 10 train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path), True) train_lineage.end(self.my_run_context(self.run_context)) args[6].assert_called() @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_train_end_exception(self, *args): """Test TrainLineage.end method when exception.""" args[0].return_value = True train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaises(Exception) as context: train_lineage.end(self.run_context) self.assertTrue('Invalid TrainLineage run_context.' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_model_size') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_file_path') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.LineageSummary.record_train_lineage') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_dataset') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_optimizer') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_network') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_train_run_context') @mock.patch('builtins.float') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_train_end_exception_log_error(self, *args): """Test TrainLineage.end method with logging errors.""" args[1].return_value = 2.0 args[2].return_value = True args[3].return_value = True args[4].return_value = None args[5].return_value = None args[6].side_effect = Exception args[7].return_value = (None, None) args[8].return_value = 10 train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaises(LineageLogError) as context: train_lineage.end(self.my_run_context(self.run_context)) self.assertTrue('End error in TrainLineage:' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_model_size') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_file_path') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.LineageSummary.record_train_lineage') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_dataset') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.analyze_optimizer') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_network') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_train_run_context') @mock.patch('builtins.float') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_train_end_exception_log_error2(self, *args): """Test TrainLineage.end method with logging errors.""" args[1].return_value = 2.0 args[2].return_value = True args[3].return_value = True args[4].return_value = None args[5].return_value = None args[6].side_effect = IOError args[7].return_value = (None, None) args[8].return_value = 10 run_context = {key: None for key in self.lineage_list} run_context['loss_fn'] = MagicMock() run_context['net_outputs'] = Tensor(0.11) train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaises(LineageLogError) as context: train_lineage.end(self.my_run_context(run_context)) self.assertTrue('End error in TrainLineage:' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_eval_exception_train_id_none(self, *args): """Test EvalLineage.end method with initialization error.""" args[0].return_value = True with self.assertRaises(MindInsightException) as context: self.my_eval_module(self.my_summary_record(self.summary_log_path), raise_exception=2) self.assertTrue('Invalid value for raise_exception.' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.' 'AnalyzeObject.analyze_dataset') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_eval_run_context') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.' 'LineageSummary.record_evaluation_lineage') def test_eval_end(self, *args): """Test EvalLineage.end method.""" args[1].return_value = True args[2].return_value = True args[3].return_value = None args[0].return_value = None eval_lineage = self.my_eval_module(self.my_summary_record(self.summary_log_path)) eval_lineage.end(self.my_run_context(self.run_context)) args[0].assert_called() @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') def test_eval_end_except_run_context(self, *args): """Test EvalLineage.end method when run_context is invalid..""" args[0].return_value = True eval_lineage = self.my_eval_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaises(Exception) as context: eval_lineage.end(self.run_context) self.assertTrue('Invalid EvalLineage run_context.' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.' 'AnalyzeObject.analyze_dataset') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_eval_run_context') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.' 'LineageSummary.record_evaluation_lineage') def test_eval_end_except_log_error(self, *args): """Test EvalLineage.end method with logging error.""" args[0].side_effect = Exception args[1].return_value = True args[2].return_value = True args[3].return_value = None eval_lineage = self.my_eval_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaises(LineageLogError) as context: eval_lineage.end(self.my_run_context(self.run_context)) self.assertTrue('End error in EvalLineage' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.' 'AnalyzeObject.analyze_dataset') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_summary_record') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.validate_eval_run_context') @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.' 'LineageSummary.record_evaluation_lineage') def test_eval_end_except_log_error2(self, *args): """Test EvalLineage.end method with logging error.""" args[0].side_effect = IOError args[1].return_value = True args[2].return_value = True args[3].return_value = None eval_lineage = self.my_eval_module(self.my_summary_record(self.summary_log_path), True) with self.assertRaises(LineageLogError) as context: eval_lineage.end(self.my_run_context(self.run_context)) self.assertTrue('End error in EvalLineage' in str(context.exception)) def test_epoch_is_zero(self, *args): """Test TrainLineage.end method.""" args[0].return_value = None run_context = self.run_context run_context['epoch_num'] = 0 with self.assertRaises(MindInsightException): train_lineage = self.my_train_module(self.my_summary_record(self.summary_log_path), True) train_lineage.end(self.my_run_context(run_context)) def tearDown(self): """Teardown.""" if os.path.exists(self.summary_log_path): try: shutil.rmtree(self.summary_log_path) except IOError: pass class TestAnalyzer(TestCase): """Test Analyzer class in model_lineage.py.""" def setUp(self): """SetUp config.""" self.analyzer = AnalyzeObject() def test_analyze_optimizer(self): """Test analyze_optimizer method.""" optimizer = Optimizer(Tensor(0.12)) res = self.analyzer.analyze_optimizer(optimizer) assert res == 0.12 def test_get_dataset_path(self): """Test get_dataset_path method.""" dataset = MindDataset( dataset_file='/path/to/mindrecord' ) res = self.analyzer.get_dataset_path(dataset) assert res == '/path/to/mindrecord' def test_get_dataset_path_wrapped(self): """Test get_dataset_path_wrapped method.""" dataset = Dataset() dataset.input.append( MindDataset( dataset_size=10, dataset_file='/path/to/cifar10' )) res = self.analyzer.get_dataset_path_wrapped(dataset) assert res == '/path/to/cifar10' @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.' 'AnalyzeObject.get_dataset_path_wrapped') def test_analyze_dataset(self, mock_get_path): """Test analyze_dataset method.""" mock_get_path.return_value = '/path/to/mindinsightset' dataset = MindDataset( dataset_size=10, dataset_file='/path/to/mindinsightset' ) res1 = self.analyzer.analyze_dataset(dataset, {'step_num': 10, 'epoch': 2}, 'train') res2 = self.analyzer.analyze_dataset(dataset, {'step_num': 5}, 'valid') assert res1 == {'step_num': 10, 'train_dataset_path': '/path/to', 'train_dataset_size': 50, 'epoch': 2} assert res2 == {'step_num': 5, 'valid_dataset_path': '/path/to', 'valid_dataset_size': 50} def test_get_dataset_path_dataset(self): """Test get_dataset_path method with Dataset.""" dataset = Dataset( dataset_size=10, dataset_path='/path/to/cifar10' ) with self.assertRaises(IndexError): self.analyzer.get_dataset_path(output_dataset=dataset) def test_get_dataset_path_mindrecord(self): """Test get_dataset_path method with MindDataset.""" dataset = MindDataset( dataset_file='/path/to/cifar10' ) dataset_path = self.analyzer.get_dataset_path(output_dataset=dataset) self.assertEqual(dataset_path, '/path/to/cifar10') def test_get_file_path(self): """Test get_file_path method.""" model_ckpt = ModelCheckpoint(prefix='', directory='/path/to') summary_step = SummaryStep(MagicMock(full_file_name='/path/to/summary.log')) list_callback = [model_ckpt, summary_step] ckpt_file_path, _ = AnalyzeObject.get_file_path(list_callback) self.assertEqual(ckpt_file_path, '/path/to/test_model.ckpt') @mock.patch('os.path.getsize') def test_get_file_size(self, os_get_size_mock): """Test get_file_size method.""" os_get_size_mock.return_value = 128 file_size = AnalyzeObject.get_file_size('/file/path') self.assertEqual(file_size, 128) @mock.patch('os.path.getsize') def test_get_file_size_except(self, os_get_size_mock): """Test failed to get the size of file.""" os_get_size_mock.side_effect = OSError analyzer = AnalyzeObject with self.assertRaises(LineageGetModelFileError) as context: analyzer.get_file_size('/file/path') self.assertTrue('Error when get model file size:' in str(context.exception)) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_file_size') def test_get_model_size(self, get_file_size_mock): """Test get_model_size method.""" get_file_size_mock.return_value = 128 analyzer = AnalyzeObject file_size = analyzer.get_model_size(ckpt_file_path='/file/path') self.assertEqual(file_size, 128) @mock.patch('mindinsight.lineagemgr.collection.model.model_lineage.AnalyzeObject.get_file_size') def test_get_model_size_no_ckpt(self, get_file_size_mock): """Test get_model_size method without ckpt file.""" get_file_size_mock.return_value = 0 analyzer = AnalyzeObject file_size = analyzer.get_model_size(ckpt_file_path='') self.assertEqual(file_size, 0) @mock.patch('builtins.vars') def test_get_optimizer_by_network(self, mock_vars): """Test get_optimizer_by_network.""" mock_optimizer = Optimizer(Tensor(0.1)) mock_cells = MagicMock() mock_cells.items.return_value = [{'key': mock_optimizer}] mock_vars.return_value = { '_cells': { 'key': mock_optimizer } } res = AnalyzeObject.get_optimizer_by_network(MagicMock()) self.assertEqual(res, mock_optimizer) @mock.patch('builtins.vars') def test_get_loss_fn_by_network(self, mock_vars): """Test get_loss_fn_by_network.""" mock_cell1 = {'_cells': {'key': SoftmaxCrossEntropyWithLogits(0.2)}} mock_cell2 = {'_cells': {'opt': Optimizer(Tensor(0.1))}} mock_cell3 = {'_cells': {'loss': SoftmaxCrossEntropyWithLogits(0.1)}} mock_vars.side_effect = [mock_cell1, mock_cell2, mock_cell3] res = AnalyzeObject.get_loss_fn_by_network(MagicMock()) self.assertEqual(res, mock_cell3['_cells']['loss']) @mock.patch('builtins.vars') def test_get_backbone_network_with_loss_cell(self, mock_vars): """Test get_backbone_network with loss_cell.""" mock_cell = {'_cells': {'key': WithLossCell(MagicMock(), SoftmaxCrossEntropyWithLogits(0.1))} } mock_vars.return_value = mock_cell res = AnalyzeObject.get_backbone_network(MagicMock()) self.assertEqual(res, 'MagicMock') @mock.patch('builtins.vars') def test_get_backbone_network(self, mock_vars): """Test get_backbone_network.""" mock_net = TrainOneStepWithLossScaleCell() mock_net.network = MagicMock() mock_cell = { '_cells': { 'key': mock_net } } mock_vars.return_value = mock_cell res = AnalyzeObject.get_backbone_network(MagicMock()) self.assertEqual(res, 'MagicMock') if __name__ == '__main__': unittest.main(verbosity=2)