mindspore 2.2.0__cp38-cp38-manylinux1_x86_64.whl → 2.2.11__cp38-cp38-manylinux1_x86_64.whl

mindspore/profiler/profiling.py

@@ -48,12 +48,12 @@ from mindspore.profiler.parser.msadvisor_analyzer import Msadvisor
 from mindspore.profiler.parser.profiler_info import ProfilerInfo
 from mindspore.common.api import _pynative_executor
 from mindspore.profiler.parser.ascend_msprof_exporter import AscendMsprofExporter
-from mindspore.profiler.parser.ascend_msprof_generator import AscendMsprofDataGenerator
+from mindspore.profiler.parser.ascend_msprof_generator import AscendMsprofDataGenerator, AscendMsprofDataGeneratorOld
 from mindspore.profiler.parser.ascend_fpbp_generator import AscendFPBPGenerator
 from mindspore.profiler.parser.ascend_op_generator import AscendOPGenerator
 from mindspore.profiler.parser.ascend_steptrace_generator import AscendStepTraceGenerator
 from mindspore.profiler.parser.ascend_flops_generator import AscendFlopsGenerator
-from mindspore.profiler.parser.ascend_hccl_generator import AscendHCCLGenerator
+from mindspore.profiler.parser.ascend_hccl_generator import AscendHCCLGenerator, AscendHCCLGeneratorOld
 
 INIT_OP_NAME = 'Default/InitDataSetQueue'
 
@@ -274,16 +274,20 @@ def _parse_host_info(input_file, output_timeline_file, output_memory_file, is_de
 
 
 def _ascend_graph_msprof_generator(source_path, model_iteration_dict):
+    """Executing the msprof export mode."""
     try:
+        ProfilerInfo.set_export_start_time(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
         msprof_exporter = AscendMsprofExporter(source_path)
-        msprof_exporter.export(model_iteration_dict)
+        flag = msprof_exporter.export(model_iteration_dict)
+        ProfilerInfo.set_export_end_time(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
+        return flag
+
     except ProfilerException as err:
         logger.warning(err.message)
-    finally:
-        pass
+        return False
 
 
-def _ascend_graph_msprof_analyse(source_path):
+def _ascend_graph_msprof_analyse(source_path, flag):
     """
     Ascend graph model msprof data analyse.
 
@@ -294,7 +298,10 @@ def _ascend_graph_msprof_analyse(source_path):
     df_op_statistic = []
     df_step_trace = []
     try:
-        msprof_analyser = AscendMsprofDataGenerator(os.path.join(source_path, 'summary'))
+        if flag:
+            msprof_analyser = AscendMsprofDataGenerator(os.path.join(source_path, 'summary'))
+        else:
+            msprof_analyser = AscendMsprofDataGeneratorOld(os.path.join(source_path, 'summary'))
         df_op_summary, df_op_statistic, df_step_trace = msprof_analyser.parse()
     except ProfilerException as err:
         logger.warning(err.message)
@@ -436,6 +443,7 @@ class Profiler:
         self._ascend_profiler = None
         self._timeline_size_limit_byte = 500 * 1024 * 1024  # 500MB
         self._parallel_strategy = True
+        self._model_iteration_dict = None
         _environment_check()
         # default aicore_metrics type is ArithmeticUtilization
         self._aicore_metrics_id = 0
@@ -450,7 +458,6 @@ class Profiler:
         self._sync_enable = True
         self._stop_time = 0
         self._dynamic_status = False
-        self._model_iteration_dict = None
         self._profile_framework = "all"
         self._msprof_enable = os.getenv("PROFILER_SAMPLECONFIG")
         if self._msprof_enable:
@@ -609,9 +616,21 @@ class Profiler:
             model_iteration_dict: Dictionary with model id as the key and iteration id as the value, Default: ``None``.
         """
         self._model_iteration_dict = model_iteration_dict
+
+        self._init_profiler_info()
+        self._is_support_step_info_collect()
+        parallel_mode = get_auto_parallel_context("parallel_mode")
+        stage_num = get_auto_parallel_context("pipeline_stages")
+
+        ProfilerInfo.set_parallel_info(parallel_mode, stage_num)
+        ProfilerInfo.set_rank_size(self._rank_size)
+        ProfilerInfo.set_heterogeneous(self._is_heterogeneous)
         if offline_path:
             if self._is_offline_parser():
+                ProfilerInfo.set_analyse_start_time(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
                 self._ascend_graph_analyse()
+                ProfilerInfo.set_analyse_end_time(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
+                ProfilerInfo.save(self._output_path)
             _offline_parse(offline_path)
             return
         if self._msprof_enable:
@@ -645,15 +664,7 @@ class Profiler:
                 logger.warning("The parameter 'profile_framework' is not support for CPU, so there no host_info"
                                " directory in the output path.")
         logger.info("Profiling: all the data have been analyzed.")
-        self._init_profiler_info()
-        self._is_support_step_info_collect()
-        parallel_mode = get_auto_parallel_context("parallel_mode")
-        stage_num = get_auto_parallel_context("pipeline_stages")
-
-        ProfilerInfo.set_parallel_info(parallel_mode, stage_num)
         ProfilerInfo.set_analyse_end_time(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
-        ProfilerInfo.set_rank_size(self._rank_size)
-        ProfilerInfo.set_heterogeneous(self._is_heterogeneous)
         ProfilerInfo.save(self._output_path)
 
     def start(self):
@@ -785,6 +796,8 @@ class Profiler:
 
             self._stop_time = int(time.time() * 10000000)
         ProfilerInfo.set_profiling_stop_time(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
+        self._init_profiler_info()
+        ProfilerInfo.save(self._output_path)
         logger.info("Profiling: stop time: %d", self._stop_time)
 
     def _profiler_init(self, kwargs):
@@ -1158,7 +1171,7 @@ class Profiler:
 
     def _ascend_flops_analyse(self, op_summary):
         """Get op FLOPs from op_summary, write output_op_flops_x.csv."""
-        if len(op_summary.dtype) != 18:
+        if 'vector_fops' not in op_summary.dtype.names and 'cube_fops' not in op_summary.dtype.names:
             logger.warning("[Profiler] Can not found cube fops and vector fops data in the op summary.")
             return
 
@@ -1195,7 +1208,7 @@ class Profiler:
         finally:
             pass
 
-    def _ascend_graph_hccl_analyse(self, source_path):
+    def _ascend_graph_hccl_analyse(self, source_path, steptrace, flag):
         """Analyse hccl profiler info."""
         if not self._profile_communication:
             return
@@ -1209,8 +1222,10 @@ class Profiler:
 
             hccl_raw_path = os.path.join(self._output_path, f'hccl_raw_{dev_id}.csv')
             hccl_raw_path = validate_and_normalize_path(hccl_raw_path)
-
-            hccl_analyse = AscendHCCLGenerator(os.path.join(source_path, 'timeline'))
+            if flag:
+                hccl_analyse = AscendHCCLGenerator(os.path.join(source_path, 'timeline'), steptrace)
+            else:
+                hccl_analyse = AscendHCCLGeneratorOld(os.path.join(source_path, 'timeline'))
             hccl_analyse.parse()
             hccl_analyse.write(hccl_raw_path)
 
@@ -1252,8 +1267,12 @@ class Profiler:
         source_path = os.path.join(self._output_path, job_id)
         self._minddata_analyse(source_path)
         if self._op_time:
-            _ascend_graph_msprof_generator(source_path, self._model_iteration_dict)
-            op_summary, op_statistic, steptrace = _ascend_graph_msprof_analyse(source_path)
+            flag = _ascend_graph_msprof_generator(source_path, self._model_iteration_dict)
+            if not flag:
+                logger.warning('Current driver package not support all export mode, use single export mode, '
+                               'this may lead to performance degradation. Suggest upgrading the driver package.')
+            ProfilerInfo.set_export_flag(flag)
+            op_summary, op_statistic, steptrace = _ascend_graph_msprof_analyse(source_path, flag)
             self._ascend_op_analyse(op_summary, op_statistic, self._dynamic_status)
             self._ascend_timeline_analyse(op_summary, steptrace)
             graph_ids = np.unique(op_summary['Model ID']).tolist()
@@ -1264,7 +1283,7 @@ class Profiler:
                 self._ascend_dynamic_net_analyse(op_summary)
             self._ascend_flops_analyse(op_summary)
             self._ascend_graph_memory_analyse(points)
-            self._ascend_graph_hccl_analyse(source_path)
+            self._ascend_graph_hccl_analyse(source_path, steptrace, flag)
             self._ascend_graph_msadvisor_analyse(job_id)
             ProfilerInfo.set_graph_ids(graph_ids)
 
@@ -1459,6 +1478,9 @@ class Profiler:
             job_id = self._ascend_job_id.rstrip('/').split('/')[-1]
             if job_id.startswith('PROF'):
                 device_dir = [dir for dir in os.listdir(self._ascend_job_id) if dir.startswith('device')]
+                info_file_path = get_file_path(os.path.join(self._ascend_job_id, device_dir[0]), "info.json")
+                training_rank_id, _ = self._parse_info_json(info_file_path)
+                self._rank_id = int(training_rank_id)
                 return os.path.join(job_id, device_dir[0])
             return job_id
 
@@ -1489,8 +1511,8 @@ class Profiler:
                                "profiler will ignore this job dir.", job_dir)
                 continue
 
-            _, training_device_id = self._parse_info_json(info_file_path)
             job_start_time = self._parse_start_log(start_file_path)
+            _, training_device_id = self._parse_info_json(info_file_path)
 
             if self._dev_id != training_device_id:
                 logger.debug("Find profiling find job path %s, but not current training device id. "

mindspore/rewrite/api/pattern_engine.py

@@ -364,7 +364,7 @@ class PatternEngine:
                 continue
             if cur_node.get_node_type() == NodeType.Tree:
                 subtree = TreeNodeHelper.get_sub_tree(cur_node)
-                self.apply(subtree)
+                _ = self.apply(subtree)
                 visited.append(cur_node)
                 queue.extend(cur_node.get_users())
                 continue

mindspore/rewrite/parsers/for_parser.py

@@ -72,7 +72,7 @@ class ForParser(Parser):
             return
         iter_code = astunparse.unparse(node.iter)
         if not iter_code.startswith(EVAL_WHITE_LIST):
-            logger.warning(
+            logger.info(
                 f"For MindSpore Rewrtie, illegal iteration condition for For node, it must start with{EVAL_WHITE_LIST}")
             return
         if "self" in iter_code:
@@ -82,7 +82,7 @@ class ForParser(Parser):
         except (NameError, TypeError) as e:
             _info = f"For MindSpore Rewrtie, when eval '{iter_code}' by using JIT Fallback feature, " \
                          f"an error occurred: {str(e)}"
-            logger.warning(_info)
+            logger.info(_info)
             stree.try_append_python_node(node, node, node_manager)
             return
 
@@ -107,7 +107,7 @@ class ForParser(Parser):
                     ast_functiondef.body.insert(index, new_node)
                     index += 1
             # Expand "for" statement and replace the body with Pass
-            for body in node.body:
+            for body in node.body[:]:
                 node.body.remove(body)
             node.body.append(ast.Pass())
 
@@ -115,13 +115,13 @@ class ForParser(Parser):
                 stree.on_change(Event.CodeChangeEvent)
             return
         if isinstance(iter_obj, range):
-            logger.warning("For MindSpore Rewrite, range not support.")
+            logger.info("For MindSpore Rewrite, range not support.")
         elif isinstance(iter_obj, zip):
-            logger.warning("For MindSpore Rewrite, zip not support.")
+            logger.info("For MindSpore Rewrite, zip not support.")
         elif isinstance(iter_obj, enumerate):
-            logger.warning("For MindSpore Rewrite, enumerate not support.")
+            logger.info("For MindSpore Rewrite, enumerate not support.")
         else:
-            logger.warning(f"For MindSpore Rewrite, not supported type: {type(iter_obj).__name__}")
+            logger.info(f"For MindSpore Rewrite, not supported type: {type(iter_obj).__name__}")
         stree.try_append_python_node(node, node, node_manager)
         return
 

mindspore/rewrite/parsers/module_parser.py

@@ -170,15 +170,15 @@ class ModuleParser(Parser):
                 level_count += 1
                 continue
             except Exception as e: # pylint: disable=W0703
-                logger.warning(f"For MindSpore Rewrite, in module parser, process import code: "
-                               f"{import_code} failed: {e}. Ignore this import code.")
+                logger.info(f"For MindSpore Rewrite, in module parser, process import code: "
+                            f"{import_code} failed: {e}. Ignore this import code.")
                 return None, None
             else:
                 # try test code success
                 return import_node_test.module, file_path_tmp
         # try codes with all level failed
-        logger.warning(f"For MindSpore Rewrite, in module parser, test import code: "
-                       f"{astunparse.unparse(import_node).strip()} failed. Ignore this import code.")
+        logger.info(f"For MindSpore Rewrite, in module parser, test import code: "
+                    f"{astunparse.unparse(import_node).strip()} failed. Ignore this import code.")
         return None, None
 
     @staticmethod

mindspore/rewrite/symbol_tree.py

@@ -226,7 +226,6 @@ class SymbolTree(Observer, Observable, NodeManager):
                 if class_str not in classes:
                     classes.add(node.name)
                     return node
-                return
 
             def visit_Try(self, node: ast.Try) -> Any:
                 if isinstance(node.body[0], (ast.Import, ast.ImportFrom)):
@@ -234,14 +233,12 @@ class SymbolTree(Observer, Observable, NodeManager):
                     if import_str not in imports:
                         imports.add(import_str)
                         return node
-                return
 
             def visit_Import(self, node: ast.Import) -> Any:
                 import_str = astunparse.unparse(node)
                 if import_str not in imports:
                     imports.add(import_str)
                     return node
-                return
 
             def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:
                 """
@@ -818,7 +815,7 @@ class SymbolTree(Observer, Observable, NodeManager):
         for node in new_nodes:
             self.insert_node(node, base_node, False, node_manager, True)
             base_node = node
-        self.erase_node(old_node)
+        _ = self.erase_node(old_node)
         return new_nodes[-1]
 
     def set_node_arg(self, node: Union[Node, str], index: int, arg: Union[ScopedValue, str]):

mindspore/run_check/_check_version.py

@@ -259,7 +259,7 @@ class AscendEnvChecker(EnvChecker):
 
     def __init__(self, library_path):
         self.library_path = library_path
-        self.version = ["7.0"]
+        self.version = ["7.1"]
         atlas_nnae_version = "/usr/local/Ascend/nnae/latest/compiler/version.info"
         atlas_toolkit_version = "/usr/local/Ascend/ascend-toolkit/latest/compiler/version.info"
         hisi_fwk_version = "/usr/local/Ascend/latest/compiler/version.info"
@@ -398,11 +398,13 @@ class AscendEnvChecker(EnvChecker):
 
     def set_env(self):
         curr_path = os.path.abspath(os.path.dirname(__file__))
+        cust_aicpu_path = os.path.abspath(os.path.join(curr_path, "../lib/plugin/ascend/custom_aicpu_ops"))
+        cust_aicore_path = os.path.abspath(os.path.join(curr_path, "../lib/plugin/ascend/custom_aicore_ops"))
         if os.getenv('ASCEND_CUSTOM_OPP_PATH'):
             os.environ['ASCEND_CUSTOM_OPP_PATH'] = os.environ['ASCEND_CUSTOM_OPP_PATH'] + ":" + \
-                                                   curr_path + "/../lib/plugin/ascend/custom_aicpu_ops"
+                                                   cust_aicore_path + ":" + cust_aicpu_path
         else:
-            os.environ['ASCEND_CUSTOM_OPP_PATH'] = curr_path + "/../lib/plugin/ascend/custom_aicpu_ops"
+            os.environ['ASCEND_CUSTOM_OPP_PATH'] = cust_aicore_path + ":" + cust_aicpu_path
         plugin_dir = os.path.dirname(self.library_path)
         akg_dir = os.path.join(plugin_dir, "ascend")
         AscendEnvChecker._concat_variable('LD_LIBRARY_PATH', akg_dir)

mindspore/safeguard/rewrite_obfuscation.py

@@ -17,13 +17,14 @@ import os
 import re
 import secrets
 from pathlib import Path
-import numpy as np
 
 from mindspore import ops, nn
 from mindspore.common.tensor import Tensor
 from mindspore import log as logger
 from mindspore import load_checkpoint, save_checkpoint
 from mindspore.rewrite import SymbolTree, Node, NodeType, TreeNodeHelper, ScopedValue
+from mindspore.rewrite.parsers.class_def_parser import ClassDefParser
+from mindspore.rewrite.parsers.class_def_parser import ModuleParser
 
 OBF_RATIOS_LENGTH = 1
 MAX_OBF_RATIOS_NUM = 50
@@ -31,7 +32,7 @@ OBF_RATIOS_WIDTH = 0
 OBF_RATIOS_INSERT_INDEX = 0
 
 
-def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./'):
+def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./', obfuscate_scale=100):
     """
     obfuscate the plaintext checkpoint files. Usually used in conjunction with
     :func:`mindspore.load_obf_params_into_net`.
@@ -49,8 +50,9 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./'):
             (such as transformer layers or resnet blocks). If target_modules is ``None``, the function would search
             target modules by itself. If found, the searched target module would be used, otherwise suggested target
             modules would be given with warning log. Default: ``None``.
-        saved_path (str): The directory path for saving obfuscated ckpt files and obf_ratios (a numpy file). obf_ratios
-            is the necessary data that needs to be load when running obfuscated network. Default: ``'./'``.
+        saved_path (str): The directory path for saving obfuscated ckpt files. Default: ``'./'``.
+        obfuscate_scale (Union[float, int]): Obfuscate scale of weights. The generated random obf_ratios will be in
+            range of (1 / obfuscate_scale, obfuscate_scale). Default: 100.
 
     Raises:
         TypeError: If `network` is not nn.Cell.
@@ -66,6 +68,9 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./'):
         ValueError: If the third string of `target_modules` is not in the format of 'obfuscate_layers:all' or
             'obfuscate_layers:int'.
 
+    Returns:
+        list[float], obf_ratios, which is the necessary data that needs to be load when running obfuscated network.
+
     Examples:
         >>> from mindspore import obfuscate_ckpt, save_checkpoint
         >>> # Refer to https://gitee.com/mindspore/docs/blob/r2.2/docs/mindspore/code/lenet.py
@@ -88,10 +93,13 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./'):
     if not _check_valid_target(network, to_split_modules):
         raise ValueError("The obfuscate module path {} is not exist, please check the input 'target_modules'."
                          .format(to_split_modules))
+    if (not isinstance(obfuscate_scale, (float, int))) or (obfuscate_scale <= 1):
+        raise ValueError("obfuscate_scale must be float or int, and larger than 1, but got {}."
+                         .format(obfuscate_scale))
     # generate and save obf_ratios to saved_path
     path_list = to_split_modules[0].split('/')
     target_list = to_split_modules[1].split('|')
-    global OBF_RATIOS_WIDTH, OBF_RATIOS_LENGTH
+    global OBF_RATIOS_LENGTH
     number_of_ratios = OBF_RATIOS_LENGTH * OBF_RATIOS_WIDTH
     if number_of_ratios > MAX_OBF_RATIOS_NUM:
         OBF_RATIOS_LENGTH = MAX_OBF_RATIOS_NUM // OBF_RATIOS_WIDTH
@@ -99,14 +107,13 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./'):
     obf_ratios = []
     secrets_generator = secrets.SystemRandom()
     for _ in range(number_of_ratios):
-        secure_float = secrets_generator.uniform(0.01, 100)
+        secure_float = secrets_generator.uniform(1 / obfuscate_scale, obfuscate_scale)
         obf_ratios.append(secure_float)
-    np.save(os.path.abspath(saved_path) + '/' + 'obf_ratios.npy', np.array(obf_ratios))
     # start obfuscate ckpt
     ckpt_dir_files = os.listdir(ckpt_files)
     for ckpt_name in ckpt_dir_files:
-        if Path(ckpt_files + ckpt_name).is_dir():
-            sub_path = ckpt_files + ckpt_name
+        sub_path = os.path.abspath(ckpt_files) + '/' + ckpt_name
+        if Path(sub_path).is_dir():
             sub_ckpt_file_list = os.listdir(sub_path)
             new_saved_path = os.path.abspath(saved_path) + '/' + ckpt_name
             if not os.path.exists(new_saved_path):
@@ -124,20 +131,24 @@ def obfuscate_ckpt(network, ckpt_files, target_modules=None, saved_path='./'):
                 continue
             _obfuscate_single_ckpt(os.path.abspath(ckpt_files) + '/' + ckpt_name, obf_ratios, path_list,
                                    target_list, saved_path)
+    return obf_ratios
 
 
 def _obfuscate_single_ckpt(ckpt_name, obf_ratios, path_list, target_list, saved_path):
     """Obfuscate single ckpt file"""
     module_has_been_obfuscated = set()
-    ckpt_param = load_checkpoint(ckpt_name)
+    try:
+        ckpt_param = load_checkpoint(ckpt_name)
+    except (ValueError, TypeError, OSError):
+        logger.error("Load checkpoint failed for file {}.".format(ckpt_name))
+        return None
     obf_ratios_index = -1
-    global OBF_RATIOS_LENGTH, OBF_RATIOS_WIDTH
     for item in ckpt_param:
         module = _get_valid_module(item, path_list, target_list)
         if module:
             layer_index = _judge_layer_index(item)
             if layer_index >= OBF_RATIOS_LENGTH:
-                break
+                continue
             if module not in module_has_been_obfuscated:
                 module_has_been_obfuscated.add(module)
                 obf_ratios_index += 1
@@ -150,9 +161,10 @@ def _obfuscate_single_ckpt(ckpt_name, obf_ratios, path_list, target_list, saved_
     ckpt_file_name = ckpt_name.split('/')[-1]
     obf_ckpt_file_name = ckpt_file_name.split('.')[0] + '_obf' + '.ckpt'
     save_checkpoint(obf_param_list, os.path.abspath(saved_path) + '/' + obf_ckpt_file_name)
+    return None
 
 
-def load_obf_params_into_net(network, target_modules, obf_ratios, **kwargs):
+def load_obf_params_into_net(network, target_modules, obf_ratios, data_parallel_num=1, **kwargs):
     """
     load obfuscate ratios into obfuscated network. Usually used in conjunction with :func:`mindspore.obfuscate_ckpt`
     interface.
@@ -166,6 +178,7 @@ def load_obf_params_into_net(network, target_modules, obf_ratios, **kwargs):
             If target_modules has the third value, it should be in the format of 'obfuscate_layers:all' or
             'obfuscate_layers:int', which represents the number of layers need to be obfuscated of duplicate layers
             (such as transformer layers or resnet blocks).
+        data_parallel_num (int): The data parallel number of parallel training. Default: 1.
         obf_ratios (Tensor): The obf ratios generated when execute :func:`mindspore.obfuscate_ckpt`.
         kwargs (dict): Configuration options dictionary.
 
@@ -211,6 +224,8 @@ def load_obf_params_into_net(network, target_modules, obf_ratios, **kwargs):
         raise ValueError("obf_ratios can not be empty.")
     if not _check_valid_target(network, target_modules):
         raise ValueError("{} is not exist, please check the input 'target_modules'.".format(target_modules))
+    if (not isinstance(data_parallel_num, int)) or (data_parallel_num <= 0):
+        raise ValueError("data_parallel_num must be positive number, but got {}.".format(data_parallel_num))
     if len(target_modules) >= 1 and target_modules[0] == '/':
         target_modules[0] = ''
     path_list = target_modules[0].split('/')
@@ -219,13 +234,13 @@ def load_obf_params_into_net(network, target_modules, obf_ratios, **kwargs):
     for _ in range(path_len):
         target_list.append([])
     target_list.append(target_modules[1].split('|'))
-    global MAX_OBF_RATIOS_NUM, OBF_RATIOS_WIDTH, OBF_RATIOS_LENGTH
+    global MAX_OBF_RATIOS_NUM, OBF_RATIOS_LENGTH
     number_of_ratios = OBF_RATIOS_LENGTH * OBF_RATIOS_WIDTH
     if number_of_ratios > MAX_OBF_RATIOS_NUM:
         OBF_RATIOS_LENGTH = MAX_OBF_RATIOS_NUM // OBF_RATIOS_WIDTH
         number_of_ratios = OBF_RATIOS_LENGTH * OBF_RATIOS_WIDTH
     MAX_OBF_RATIOS_NUM = number_of_ratios
-    rewrite_network = _obfuscate_network(network, path_list, target_list, **kwargs)
+    rewrite_network = _obfuscate_network(network, path_list, target_list, data_parallel_num=data_parallel_num, **kwargs)
     setattr(rewrite_network, 'obf_ratios', obf_ratios)
     return rewrite_network
 
@@ -263,7 +278,7 @@ def _check_valid_target(network, target_modules):
     if not target_modules[1]:
         raise ValueError("{} should be a non-empty string value, in the form of 'D1|D2'"
                          .format(target_modules[1]))
-    if not re.fullmatch(pattern=r'([a-zA-Z]*[0-9]*\/*_*)*', string=target_modules[0])\
+    if not re.fullmatch(pattern=r'([a-zA-Z]*[0-9]*\/*_*)*', string=target_modules[0]) \
             or not re.fullmatch(pattern=r'([a-zA-Z]*[0-9]*\|*_*)*', string=target_modules[1]):
         raise ValueError("please check the input 'target_modules'{},it should be in the form of ['A/B/C', 'D1|D2']."
                          "target_modules[0] can only contain uppercase and lowercase letters, numbers, '_' and '/',"
@@ -297,9 +312,9 @@ def _check_valid_target(network, target_modules):
     # check whether target_list is valid
     global OBF_RATIOS_WIDTH
     OBF_RATIOS_WIDTH = 0
-    for j in range(len(target_list)):
-        if not hasattr(net, target_list[j]):
-            logger.warning("{} does not exist in the path {}".format(target_list[j], target_modules[0]))
+    for target in target_list:
+        if not hasattr(net, target):
+            logger.warning("{} does not exist in the path {}".format(target, target_modules[0]))
         else:
             OBF_RATIOS_WIDTH += 1
     if OBF_RATIOS_WIDTH == 0:
@@ -328,6 +343,7 @@ def _update_max_obf_ratios_num(target_modules):
 
 def _get_default_target_modules(ckpt_files):
     """Get the default or suggested target modules, if the target modules is None."""
+
     def _split_to_path_and_target(module, target):
         # split module into path list and target list
         target_index = module.index(target)
@@ -370,7 +386,11 @@ def _get_default_target_modules(ckpt_files):
     for ckpt_name in ckpt_dir_files:
         if not ckpt_name.endswith('.ckpt'):
             continue
-        ckpt_param = load_checkpoint(os.path.abspath(ckpt_files) + '/' + ckpt_name)
+        try:
+            ckpt_param = load_checkpoint(os.path.abspath(ckpt_files) + '/' + ckpt_name)
+        except (ValueError, TypeError, OSError):
+            logger.error("Load checkpoint failed for file {}.".format(os.path.abspath(ckpt_files) + '/' + ckpt_name))
+            return None
         for item in ckpt_param:
             param_path = _remove_digit(item)
             param_path = '/'.join(param_path)
@@ -396,9 +416,9 @@ def _get_valid_module(item, path_list, target_list):
     tar_path = '/'.join(path_list)
     # update the weights with obf_ratios in target module
     if net_path == tar_path:
-        for i in range(len(target_list)):
-            if target_list[i] in item.split('.'):
-                target_index = item.split('.').index(target_list[i])
+        for target in target_list:
+            if target in item.split('.'):
+                target_index = item.split('.').index(target)
                 module = ''.join(item.split('.')[:target_index + 1])
                 return module
     return None
@@ -413,7 +433,7 @@ def _remove_digit(item):
     return param_path
 
 
-def _obfuscate_network(model, path_list, target_list, **kwargs):
+def _obfuscate_network(model, path_list, target_list, data_parallel_num=1, **kwargs):
     """obfuscate original network, including add mul operation and add inputs for passing obf_ratio."""
 
     def _insert_input(stree: SymbolTree, arg_name: str = 'y_obf'):
@@ -426,7 +446,6 @@ def _obfuscate_network(model, path_list, target_list, **kwargs):
         # the insert input node name would be 'input_y_obf'
         new_input_node = last_input.create_input(arg_name)
         stree.insert(position, new_input_node)
-        return new_input_node
 
     def _insert_mul(stree: SymbolTree, node: Node, index: int):
         """add mul operation for original network"""
@@ -436,7 +455,12 @@ def _obfuscate_network(model, path_list, target_list, **kwargs):
         sv: ScopedValue = ScopedValue.create_naming_value(v + f'[{index}]')
         arg_list.append(sv)
         target_list = node.get_targets().copy()
-        new_mul_node = node.create_call_cell(cell=ops.Mul(), targets=target_list, args=arg_list, name='mul')
+        if data_parallel_num > 1:
+            logger.info("Data parallel number is: {}".format(data_parallel_num))
+            new_mul_node = node.create_call_cell(cell=ops.Mul().shard(((data_parallel_num, 1), ())),
+                                                 targets=target_list, args=arg_list, name='mul')
+        else:
+            new_mul_node = node.create_call_cell(cell=ops.Mul(), targets=target_list, args=arg_list, name='mul')
         position = stree.after(node)
         stree.insert(position, new_mul_node)
 
@@ -459,6 +483,8 @@ def _obfuscate_network(model, path_list, target_list, **kwargs):
         if input_y_node is None:
             return
         arg_list = subnode.get_args().copy()
+        kwargs_list = list(subnode.get_kwargs().values())
+        arg_list.extend(kwargs_list)
         v: str = input_y_node.get_targets()[0].value
         arg_obf: ScopedValue = ScopedValue.create_naming_value("y_obf=" + v)
         arg_list.append(arg_obf)
@@ -482,13 +508,11 @@ def _obfuscate_network(model, path_list, target_list, **kwargs):
 
     def _register_denied_func_decorators(fn):
         """set the function decorators which should be denied for parse"""
-        from mindspore.rewrite.parsers.class_def_parser import ClassDefParser
         name = "denied_function_decorator_list"
         setattr(ClassDefParser, name, fn)
 
     def _register_denied_class_decorators(fn):
         """set the class decorators which should be denied for parse"""
-        from mindspore.rewrite.parsers.class_def_parser import ModuleParser
         name = "denied_class_decorator_list"
         setattr(ModuleParser, name, fn)
 

mindspore/scipy/ops.py

@@ -1,4 +1,4 @@
-# Copyright 2021 Huawei Technologies Co., Ltd
+# Copyright 2021-2023 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.
@@ -156,14 +156,64 @@ class LU(PrimitiveWithInfer):
 
 
 class LinearSumAssignment(Primitive):
-    """Solve the linear sum assignment problem."""
+    r"""
+    Solve the linear sum assignment problem.
+
+    The assignment problem is represented as follows:
+
+    .. math::
+        min\sum_{i}^{} \sum_{j}^{} C_{i,j} X_{i,j}
+
+    where :math:`C` is cost matrix, :math:`X_{i,j} = 1` means column :math:`j` is assigned to row :math:`i` .
+
+    Inputs:
+        - **cost_matrix** (Tensor) - 2-D cost matrix. Tensor of shape :math:`(M, N)` .
+        - **dimension_limit** (Tensor, optional) - A scalar used to limit the actual size of the 2nd dimension of
+          ``cost_matrix``. Default is ``Tensor(sys.maxsize)``, which means no limitation. The type is 0-D int64
+          Tensor.
+        - **maximize** (bool) - Calculate a maximum weight matching if true, otherwise calculate a minimum weight
+          matching.
+
+    Outputs:
+        A tuple of tensors containing 'row_idx' and 'col_idx'.
+
+        - **row_idx** (Tensor) - Row indices of the problem. If `dimension_limit` is given, -1 would be padded at the
+          end. The shape is  :math:`(N, )` , where :math:`N` is the minimum value of `cost_matrix` dimension.
+        - **col_idx** (Tensor) - Column indices of the problem. If `dimension_limit` is given, -1 would be padded at
+          the end. The shape is  :math:`(N, )` , where :math:`N` is the minimum value of `cost_matrix` dimension.
+
+     Raises:
+        TypeError: If the data type of `cost_matrix` is not the type in [float16, float32, float64,
+                   int8, int16, int32, int64, uint8, uint16, uint32, uint64, bool]
+        TypeError: If the type of `maximize` is not bool.
+        TypeError: If the data type of `dimension_limit` is not int64.
+        ValueError: If the rank of `cost_matrix` is not 2.
+        ValueError: If the number of input args is not 3.
+
+
+    Supported Platforms:
+        ``Ascend`` ``CPU``
+
+    Examples:
+        >>> import mindspore as ms
+        >>> import numpy as np
+        >>> from mindspore import Tensor
+        >>> from mindspore.scipy.ops import LinearSumAssignment
+        >>> lsap = LinearSumAssignment()
+        >>> cost_matrix = Tensor(np.array([[2, 3, 3], [3, 2, 3], [3, 3, 2]])).astype(ms.float64)
+        >>> dimension_limit = Tensor(2)
+        >>> maximize = False
+        >>> a, b = lsap(cost_matrix, dimension_limit, maximize)
+        >>> print(a)
+        [0 1 -1]
+        >>> print(b)
+        [0 1 -1]
+    """
 
     @prim_attr_register
     def __init__(self):
-        super().__init__("LinearSumAssignment")
+        super().__init__(name="LinearSumAssignment")
         self.init_prim_io_names(inputs=['cost_matrix', 'dimension_limit', 'maximize'], outputs=['row_ind', 'col_ind'])
-        self.add_prim_attr("cust_aicpu", "mindspore_aicpu_kernels")
-
 
 # pylint: disable=C0413,W0611
 from .ops_grad import get_bprpo_eigh, get_bprpo_trsm