diff --git a/mindinsight/lineagemgr/collection/model/model_lineage.py b/mindinsight/lineagemgr/collection/model/model_lineage.py index eeb5b89cf9751431259a940e26ebf5caa7d288d1..966a59b7cc5a42e197fd79c5b210e9c818bb5fd1 100644 --- a/mindinsight/lineagemgr/collection/model/model_lineage.py +++ b/mindinsight/lineagemgr/collection/model/model_lineage.py @@ -37,7 +37,7 @@ from mindinsight.lineagemgr.collection.model.base import Metadata try: from mindspore.common.tensor import Tensor from mindspore.train.callback import Callback, RunContext, ModelCheckpoint, SummaryStep - from mindspore.nn import Cell, Optimizer, WithLossCell, TrainOneStepWithLossScaleCell + from mindspore.nn import Cell, Optimizer from mindspore.nn.loss.loss import _Loss from mindspore.dataset.engine import Dataset, MindDataset import mindspore.dataset as ds @@ -412,27 +412,28 @@ class AnalyzeObject: Returns: str, the name of the backbone network. """ - with_loss_cell = False - backbone = None + backbone_name = None + has_network = False + network_key = 'network' + backbone_key = '_backbone' + net_args = vars(network) if network else {} net_cell = net_args.get('_cells') if net_args else {} - for _, value in net_cell.items(): - if isinstance(value, WithLossCell): - backbone = getattr(value, '_backbone') - with_loss_cell = True + for key, value in net_cell.items(): + if key == network_key: + network = value + has_network = True break - if with_loss_cell: - backbone_name = type(backbone).__name__ \ - if backbone else None - elif isinstance(network, TrainOneStepWithLossScaleCell): - backbone = getattr(network, 'network') - backbone_name = type(backbone).__name__ \ - if backbone else None - else: - backbone_name = type(network).__name__ \ - if network else None + if has_network: + while hasattr(network, network_key): + network = getattr(network, network_key) + if hasattr(network, backbone_key): + backbone = getattr(network, backbone_key) + backbone_name = type(backbone).__name__ + elif network is not None: + backbone_name = type(network).__name__ return backbone_name @staticmethod @@ -489,26 +490,24 @@ class AnalyzeObject: Returns: dict, the lineage metadata. """ - dataset_batch_size = dataset.get_dataset_size() - if dataset_batch_size is not None: - validate_int_params(dataset_batch_size, 'dataset_batch_size') - log.debug('dataset_batch_size: %d', dataset_batch_size) + batch_num = dataset.get_dataset_size() + batch_size = dataset.get_batch_size() + if batch_num is not None: + validate_int_params(batch_num, 'dataset_batch_num') + validate_int_params(batch_num, 'dataset_batch_size') + log.debug('dataset_batch_num: %d', batch_num) + log.debug('dataset_batch_size: %d', batch_size) dataset_path = AnalyzeObject.get_dataset_path_wrapped(dataset) if dataset_path: dataset_path = '/'.join(dataset_path.split('/')[:-1]) - step_num = lineage_dict.get('step_num') - validate_int_params(step_num, 'step_num') - log.debug('step_num: %d', step_num) - + dataset_size = int(batch_num * batch_size) if dataset_type == 'train': lineage_dict[Metadata.train_dataset_path] = dataset_path - epoch = lineage_dict.get('epoch') - train_dataset_size = dataset_batch_size * (step_num / epoch) - lineage_dict[Metadata.train_dataset_size] = int(train_dataset_size) + lineage_dict[Metadata.train_dataset_size] = dataset_size elif dataset_type == 'valid': lineage_dict[Metadata.valid_dataset_path] = dataset_path - lineage_dict[Metadata.valid_dataset_size] = dataset_batch_size * step_num + lineage_dict[Metadata.valid_dataset_size] = dataset_size return lineage_dict diff --git a/mindinsight/lineagemgr/querier/query_model.py b/mindinsight/lineagemgr/querier/query_model.py index 74051c9996d8ca692671d2511ed8597d4948cbd9..007c2d4c83592726bb259f3ee574712457655585 100644 --- a/mindinsight/lineagemgr/querier/query_model.py +++ b/mindinsight/lineagemgr/querier/query_model.py @@ -82,7 +82,6 @@ class LineageObj: self._lineage_info = { self._name_summary_dir: summary_dir } - self._filtration_result = None self._init_lineage() self.parse_and_update_lineage(**kwargs) diff --git a/tests/st/func/lineagemgr/api/test_model_api.py b/tests/st/func/lineagemgr/api/test_model_api.py index 2357c44097ca2130c4a42c81072c48e6d7d37fd8..c8aea9dbabbb0e27adbcbbd574cdfb5600e990a8 100644 --- a/tests/st/func/lineagemgr/api/test_model_api.py +++ b/tests/st/func/lineagemgr/api/test_model_api.py @@ -50,10 +50,10 @@ LINEAGE_INFO_RUN1 = { 'network': 'ResNet' }, 'train_dataset': { - 'train_dataset_size': 731 + 'train_dataset_size': 1024 }, 'valid_dataset': { - 'valid_dataset_size': 10240 + 'valid_dataset_size': 1024 }, 'model': { 'path': '{"ckpt": "' @@ -89,9 +89,9 @@ LINEAGE_FILTRATION_RUN1 = { 'model_lineage': { 'loss_function': 'SoftmaxCrossEntropyWithLogits', 'train_dataset_path': None, - 'train_dataset_count': 731, + 'train_dataset_count': 1024, 'test_dataset_path': None, - 'test_dataset_count': 10240, + 'test_dataset_count': 1024, 'user_defined': {}, 'network': 'ResNet', 'optimizer': 'Momentum', @@ -115,7 +115,7 @@ LINEAGE_FILTRATION_RUN2 = { 'train_dataset_path': None, 'train_dataset_count': 1024, 'test_dataset_path': None, - 'test_dataset_count': 10240, + 'test_dataset_count': 1024, 'user_defined': {}, 'network': "ResNet", 'optimizer': "Momentum", diff --git a/tests/ut/lineagemgr/collection/model/test_model_lineage.py b/tests/ut/lineagemgr/collection/model/test_model_lineage.py index e6857f72d2bd40f830ec8187b6398d0f79f535d8..81a61ea6dfe4981d943d3ea108ec0e5108c05483 100644 --- a/tests/ut/lineagemgr/collection/model/test_model_lineage.py +++ b/tests/ut/lineagemgr/collection/model/test_model_lineage.py @@ -334,12 +334,14 @@ class TestAnalyzer(TestCase): ) res1 = self.analyzer.analyze_dataset(dataset, {'step_num': 10, 'epoch': 2}, 'train') res2 = self.analyzer.analyze_dataset(dataset, {'step_num': 5}, 'valid') + + # batch_size is mocked as 32. assert res1 == {'step_num': 10, 'train_dataset_path': '/path/to', - 'train_dataset_size': 50, + 'train_dataset_size': 320, 'epoch': 2} assert res2 == {'step_num': 5, 'valid_dataset_path': '/path/to', - 'valid_dataset_size': 50} + 'valid_dataset_size': 320} def test_get_dataset_path_dataset(self): """Test get_dataset_path method with Dataset.""" diff --git a/tests/utils/mindspore/dataset/engine/datasets.py b/tests/utils/mindspore/dataset/engine/datasets.py index a9decbceb9814c5b5598a501bf1990d28ff9e8dc..8fc28967d2dfdf79bf44a9a18342566823e32fb0 100644 --- a/tests/utils/mindspore/dataset/engine/datasets.py +++ b/tests/utils/mindspore/dataset/engine/datasets.py @@ -27,6 +27,10 @@ class Dataset: """Mocked get_dataset_size.""" return self.dataset_size + def get_batch_size(self): + """Mocked get_batch_size""" + return 32 + class MindDataset(Dataset): """Mock the MindSpore MindDataset class."""