mindstudio-probe 1.0.4__py3-none-any.whl → 1.1.1__py3-none-any.whl

msprobe/mindspore/free_benchmark/perturbation/add_noise.py

@@ -1,11 +1,26 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from typing import Any
 
 from mindspore import Tensor, ops
 
+from msprobe.mindspore.common.const import FreeBenchmarkConst
 from msprobe.mindspore.common.log import logger
-from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
-from msprobe.mindspore.common.const import FreeBenchmarkConst
+from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 
 
 class AddNoisePerturbation(BasePerturbation):
@@ -17,7 +32,7 @@ class AddNoisePerturbation(BasePerturbation):
         """
         params.fuzzed_value = self.add_noise(params.args[params.index])
         if not self.is_fuzzed:
-            logger.warning(f"{self.api_name} can not add noise.")
+            logger.warning(f"{self.api_name_with_id} can not add noise.")
             return False
         return self.get_fuzzed_result(params)
 
@@ -43,25 +58,25 @@ class AddNoisePerturbation(BasePerturbation):
 
         return inputs
 
-    def _get_noise(self, input):
+    def _get_noise(self, tensor):
         """
         得到要添加的噪声值
 
         """
         if self.is_fuzzed:
             return False
-        if not ops.is_floating_point(input) or ops.numel(input) == 0:
+        if not ops.is_floating_point(tensor) or ops.numel(tensor) == 0:
             return False
 
-        pert_value = FreeBenchmarkConst.PERT_VALUE_DICT.get(input.dtype)
+        pert_value = FreeBenchmarkConst.PERT_VALUE_DICT.get(tensor.dtype)
         if not pert_value:
             return False
         else:
             self.perturbation_value = pert_value
 
-        max_val = ops.max(ops.abs(input))[0].item()
+        max_val = ops.max(ops.abs(tensor))[0].item()
         if max_val < pert_value:
             return False
 
-        noise = ops.full(input.shape, self.perturbation_value, dtype=input.dtype)
+        noise = ops.full(tensor.shape, self.perturbation_value, dtype=tensor.dtype)
         return noise

msprobe/mindspore/free_benchmark/perturbation/base_perturbation.py

@@ -1,20 +1,44 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from typing import Any
 
+from msprobe.core.common.const import Const
+from msprobe.mindspore.free_benchmark.common.config import Config
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
+from msprobe.mindspore.free_benchmark.common.utils import Tools
 
 
 class BasePerturbation:
 
-    def __init__(self, api_name: str):
-        self.api_name = api_name
+    def __init__(self, api_name_with_id: str):
+        self.api_name_with_id = api_name_with_id
         self.is_fuzzed = False
         self.perturbation_value = None
 
     @staticmethod
     def get_fuzzed_result(params: HandlerParams):
-        args_front = params.args[:params.index]
-        args_rear = params.args[params.index + 1:]
-        fuzzed_result = params.original_func(*args_front, params.fuzzed_value, *args_rear, **params.kwargs)
+        if Config.stage == Const.BACKWARD:
+            fuzzed_result = Tools.get_grad(params.original_func, *params.args[:params.index],
+                                           params.fuzzed_value, *params.args[params.index + 1:], **params.kwargs)
+
+            if fuzzed_result is None:
+                return False
+        else:
+            fuzzed_result = params.original_func(*params.args[:params.index], params.fuzzed_value,
+                                                 *params.args[params.index + 1:], **params.kwargs)
         return fuzzed_result
 
     def handler(self, params: HandlerParams) -> Any:

msprobe/mindspore/free_benchmark/perturbation/bit_noise.py

@@ -1,10 +1,25 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from typing import Any
 
 import numpy as np
 from mindspore import Tensor, ops
 
-from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.common.const import FreeBenchmarkConst
+from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
 from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 
@@ -20,12 +35,12 @@ class BitNoisePerturbation(BasePerturbation):
                 noise_type = list(FreeBenchmarkConst.MS_NUMPY_DTYPE_DICT.keys())[
                              list(FreeBenchmarkConst.MS_NUMPY_DTYPE_DICT.values()).index(bit_len_type)]
                 noise = ops.full(inputs.shape, 1, dtype=noise_type)
-                input_np = inputs.asnumpy()
+                input_np = inputs.contiguous().asnumpy()
                 input_np_int = input_np.view(bit_len_type)
                 result = Tensor(input_np_int)
                 result = ops.where(ops.abs(inputs) > sub_normal,
                                    ops.bitwise_xor(result, noise), result)
-                result_np = result.asnumpy()
+                result_np = result.contiguous().asnumpy()
                 result_np_float = result_np.view(FreeBenchmarkConst.MS_NUMPY_DTYPE_DICT.get(inputs.dtype))
                 self.is_fuzzed = True
                 return Tensor(result_np_float)
@@ -40,24 +55,24 @@ class BitNoisePerturbation(BasePerturbation):
         args = params.args
         params.fuzzed_value = self.add_bit_noise(params.args[params.index])
         if not self.is_fuzzed:
-            logger.warning(f"{self.api_name} can not add bit noise.")
+            logger.warning(f"{self.api_name_with_id} can not add bit noise.")
             return False
         params.args = args
         return self.get_fuzzed_result(params)
 
-    def _get_bit_len_type(self, input):
+    def _get_bit_len_type(self, tensor):
         if self.is_fuzzed:
             return False
-        if not isinstance(input, Tensor) or not ops.is_floating_point(input) or \
-           input.numel() == 0:
+        if not isinstance(tensor, Tensor) or not ops.is_floating_point(tensor) or \
+           tensor.numel() == 0:
             return False
-        bit_len_type = FreeBenchmarkConst.PERT_BIT_DICT.get(input.dtype)
+        bit_len_type = FreeBenchmarkConst.PERT_BIT_DICT.get(tensor.dtype)
         if not bit_len_type:
             return False
-        pert_value = FreeBenchmarkConst.PERT_VALUE_DICT.get(input.dtype)
+        pert_value = FreeBenchmarkConst.PERT_VALUE_DICT.get(tensor.dtype)
         if not pert_value:
             return False
-        max_val = ops.max(ops.abs(input))[0].item()
+        max_val = ops.max(ops.abs(tensor))[0].item()
         if max_val < pert_value:
             return False
         return bit_len_type

msprobe/mindspore/free_benchmark/perturbation/exchange_value.py

@@ -1,14 +1,39 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from typing import Any
 
-from mindspore import Tensor
+from mindspore import Tensor, ops
 
 from msprobe.mindspore.common.log import logger
-from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
+from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 
 
 class ExchangeValuePerturbation(BasePerturbation):
 
+    @staticmethod
+    def _check_tensor_shape(inputs):
+        dims = len(inputs.shape)
+        if dims == 1 and inputs.shape[0] > 1:
+            return True
+        if dims > 1 and inputs.shape[1] > 0:
+            if inputs.shape[0] > 1 or inputs.shape[1] > 1:
+                return True
+        return False
+
     def handle(self, params: HandlerParams) -> Any:
         """
         返回首尾交换后的api输出
@@ -16,7 +41,7 @@ class ExchangeValuePerturbation(BasePerturbation):
         """
         params.fuzzed_value = self.exchange_value(params.args[params.index])
         if not self.is_fuzzed:
-            logger.warning(f"{self.api_name} can not exchange value.")
+            logger.warning(f"{self.api_name_with_id} can not exchange value.")
             return False
         return self.get_fuzzed_result(params)
 
@@ -25,22 +50,23 @@ class ExchangeValuePerturbation(BasePerturbation):
         返回首尾交换后的api输入
 
         """
-        if isinstance(inputs, Tensor):
-            if not self.is_fuzzed and len(inputs.shape) > 0 and inputs.shape[0] > 1:
-                result = inputs.copy()
-                if len(inputs.shape) == 1:
-                    first_element = inputs[0]
-                    last_element = inputs[-1]
-                    result[0] = last_element
-                    result[-1] = first_element
-                else:
-                    first_element = inputs[0][0]
-                    last_element = inputs[-1][-1]
-                    result[0][0] = last_element
-                    result[-1][-1] = first_element
-
-                self.is_fuzzed = True
-                return result
+        if isinstance(inputs, Tensor) and ops.is_floating_point(inputs):
+            if self.is_fuzzed or not self._check_tensor_shape(inputs):
+                return inputs
+            result = inputs.copy()
+            if len(inputs.shape) == 1:
+                first_element = inputs[0]
+                last_element = inputs[-1]
+                result[0] = last_element
+                result[-1] = first_element
+            else:
+                first_element = inputs[0][0]
+                last_element = inputs[-1][-1]
+                result[0][0] = last_element
+                result[-1][-1] = first_element
+
+            self.is_fuzzed = True
+            return result
 
         if isinstance(inputs, dict):
             return {k: self.exchange_value(v) for k, v in inputs.items()}

msprobe/mindspore/free_benchmark/perturbation/improve_precision.py

@@ -1,13 +1,29 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from typing import Any
 
 import mindspore as ms
 from mindspore import Tensor, ops
 
-from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
-from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
-from msprobe.mindspore.common.const import FreeBenchmarkConst
+from msprobe.core.common.const import Const
 from msprobe.mindspore.common.log import logger
-from msprobe.mindspore.common.const import Const
+from msprobe.mindspore.free_benchmark.common.config import Config
+from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
+from msprobe.mindspore.free_benchmark.common.utils import Tools
+from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 
 
 class ImprovePrecisionPerturbation(BasePerturbation):
@@ -26,10 +42,15 @@ class ImprovePrecisionPerturbation(BasePerturbation):
     def handle(self, params: HandlerParams) -> Any:
         args = self.improve_tensor_precision(params.args)
         kwargs = self.improve_tensor_precision(params.kwargs)
-        fuzzed_value = args
-        if self.api_name in Const.COMMUNICATION_API_LIST:
-            params.fuzzed_value = fuzzed_value
         if not self.is_fuzzed:
-            logger.warning(f"{self.api_name} can not improve precision.")
+            logger.warning(f"{self.api_name_with_id} can not improve precision.")
             return False
+
+        if Config.stage == Const.BACKWARD:
+            fuzzed_result = Tools.get_grad(params.original_func, *args, **kwargs)
+            if fuzzed_result is not None:
+                return fuzzed_result
+            else:
+                return False
+
         return params.original_func(*args, **kwargs)

msprobe/mindspore/free_benchmark/perturbation/no_change.py

@@ -1,7 +1,22 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from typing import Any
 
-from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 from msprobe.mindspore.free_benchmark.common.handler_params import HandlerParams
+from msprobe.mindspore.free_benchmark.perturbation.base_perturbation import BasePerturbation
 
 
 class NoChangePerturbation(BasePerturbation):

msprobe/mindspore/free_benchmark/perturbation/perturbation_factory.py

@@ -1,10 +1,25 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from msprobe.mindspore.common.const import FreeBenchmarkConst
 from msprobe.mindspore.free_benchmark.common.config import Config
-from .add_noise import AddNoisePerturbation
-from .bit_noise import BitNoisePerturbation
-from .no_change import NoChangePerturbation
-from .improve_precision import ImprovePrecisionPerturbation
-from .exchange_value import ExchangeValuePerturbation
+from msprobe.mindspore.free_benchmark.perturbation.add_noise import AddNoisePerturbation
+from msprobe.mindspore.free_benchmark.perturbation.bit_noise import BitNoisePerturbation
+from msprobe.mindspore.free_benchmark.perturbation.exchange_value import ExchangeValuePerturbation
+from msprobe.mindspore.free_benchmark.perturbation.improve_precision import ImprovePrecisionPerturbation
+from msprobe.mindspore.free_benchmark.perturbation.no_change import NoChangePerturbation
 
 
 class PerturbationFactory:
@@ -21,9 +36,9 @@ class PerturbationFactory:
     }
 
     @staticmethod
-    def create(api_name: str):
+    def create(api_name_with_id: str):
         perturbation = PerturbationFactory.perturbations.get(Config.pert_type)
         if perturbation:
-            return perturbation(api_name)
+            return perturbation(api_name_with_id)
         else:
             raise Exception(f'{Config.pert_type} is a invalid perturbation type')

msprobe/mindspore/free_benchmark/self_check_tool_factory.py

@@ -1,6 +1,21 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 from msprobe.mindspore.common.const import Const
 from msprobe.mindspore.debugger.debugger_config import DebuggerConfig
-from msprobe.mindspore.free_benchmark.api_pynative_self_check import ApiPyNativeSelFCheck
+from msprobe.mindspore.free_benchmark.api_pynative_self_check import ApiPyNativeSelfCheck
 
 
 class SelfCheckToolFactory:
@@ -13,7 +28,7 @@ class SelfCheckToolFactory:
         Const.API: {
             Const.GRAPH_KBYK_MODE: None,
             Const.GRAPH_GE_MODE: None,
-            Const.PYNATIVE_MODE: ApiPyNativeSelFCheck
+            Const.PYNATIVE_MODE: ApiPyNativeSelfCheck
         },
         Const.KERNEL: {
             Const.GRAPH_KBYK_MODE: None,

msprobe/mindspore/grad_probe/global_context.py

@@ -1,15 +1,30 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
 import os
 import threading
-from typing import Dict, Union
+from typing import Dict, Union, Tuple
 
-from msprobe.core.grad_probe.utils import check_str
+from msprobe.core.common.utils import is_int
+from msprobe.core.common.file_utils import create_directory, check_path_before_create
 from msprobe.core.grad_probe.constant import GradConst
+from msprobe.core.grad_probe.utils import check_str, check_bounds_element, check_param_element
 from msprobe.mindspore.common.log import logger
-from msprobe.core.common.file_utils import create_directory, check_path_before_create
 
 
 class GlobalContext:
-
     _instance = None
     _instance_lock = threading.Lock()
     _setting = {
@@ -18,7 +33,7 @@ class GlobalContext:
         GradConst.STEP: None,
         GradConst.RANK: None,
         GradConst.CURRENT_STEP: 0,
-        GradConst.BOUNDS: [-10, -1, -0.1, -0.01, -0.001, 0, 0.001, 0.01, 0.1, 1, 10],
+        GradConst.BOUNDS: [-1, 0, 1],
         GradConst.OUTPUT_PATH: None
     }
 
@@ -31,19 +46,19 @@ class GlobalContext:
 
     def init_context(self, config_dict: Dict):
         level = config_dict.get(GradConst.LEVEL)
-        check_str(level, variable_name = "level in yaml")
+        check_str(level, variable_name="level in yaml")
         if level in GradConst.SUPPORTED_LEVEL:
             self._setting[GradConst.LEVEL] = config_dict.get(GradConst.LEVEL)
         else:
             raise ValueError("Invalid level set in config yaml file, level option: L0, L1, L2")
 
-        self._set_input_list(config_dict, GradConst.PARAM_LIST, str)
-        self._set_input_list(config_dict, GradConst.BOUNDS, float)
-        self._set_input_list(config_dict, GradConst.STEP, int)
-        self._set_input_list(config_dict, GradConst.RANK, int)
+        self._set_input_list(config_dict, GradConst.PARAM_LIST, (str,), element_check=check_param_element)
+        self._set_input_list(config_dict, GradConst.BOUNDS, (float, int), element_check=check_bounds_element)
+        self._set_input_list(config_dict, GradConst.STEP, (int,))
+        self._set_input_list(config_dict, GradConst.RANK, (int,))
 
         output_path = config_dict.get(GradConst.OUTPUT_PATH)
-        check_str(output_path, variable_name = "output_path in yaml")
+        check_str(output_path, variable_name="output_path in yaml")
         try:
             check_path_before_create(output_path)
         except RuntimeError as err:
@@ -70,21 +85,36 @@ class GlobalContext:
         dump_rank_list = self.get_context(GradConst.RANK)
         return (not dump_rank_list) or (rank in dump_rank_list)
 
-    def _set_input_list(self, config_dict: Dict, name: str, dtype: Union[int, str, float]):
-        value = config_dict.get(name)
+    def _get_type_str(self, dtype: Union[int, str, float, Tuple[int, str, float]]):
+        if isinstance(dtype, tuple):
+            return "/".join([self._get_type_str(element) for element in dtype])
         if dtype == int:
             type_str = "integer"
         elif dtype == float:
             type_str = "float"
         else:
             type_str = "string"
+        return type_str
+
+    def _set_input_list(self, config_dict: Dict, name: str,
+                        dtype: Union[int, str, float, Tuple[int, str, float]], element_check=None):
+        value = config_dict.get(name)
+        type_str = self._get_type_str(dtype)
         if value and isinstance(value, list):
             for val in value:
                 if not isinstance(val, dtype):
-                    logger.warning(f"Invalid {name} which must be None or list of {type_str}")
+                    logger.warning(f"Invalid {name} which must be None or list of {type_str}, use default value.")
+                    return
+                elif isinstance(val, int) and not is_int(val):
+                    logger.warning(f"Invalid {name} which must be None or list of int, use default value.")
+                    return
+                if element_check and not element_check(val):
+                    logger.warning(f"Given {name} violates some rules, use default value.")
                     return
+
             self._setting[name] = value
         else:
             logger.warning(f"{name} is None or not a list with valid items, use default value.")
 
+
 grad_context = GlobalContext()

msprobe/mindspore/grad_probe/grad_analyzer.py

@@ -1,20 +1,33 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
+import multiprocessing
 import os
 import time
-from typing import List, Tuple
-import multiprocessing
 from multiprocessing import Process
+from typing import List
 
-import numpy as np
 import mindspore as ms
-from mindspore.communication import get_rank
-from mindspore.ops import operations as P
+import numpy as np
 from mindspore.common.parameter import Parameter
-
-from msprobe.core.grad_probe.utils import ListCache
-from msprobe.core.grad_probe.constant import GradConst
-from msprobe.mindspore.common.log import logger
+from mindspore.communication import get_rank
 from msprobe.core.common.file_utils import (create_directory, check_file_or_directory_path,
                                             write_csv, remove_path, move_file, load_npy)
+from msprobe.core.grad_probe.constant import GradConst
+from msprobe.core.grad_probe.utils import ListCache
+from msprobe.mindspore.common.log import logger
 from msprobe.mindspore.grad_probe.global_context import grad_context, GlobalContext
 
 
@@ -28,12 +41,12 @@ def get_rank_id():
 
 @ms.jit
 def grad_dump(dump_dir: str, g_name: str, dump_step: Parameter, grad: ms.Tensor, level: str, bounds: List):
-    '''
+    """
     Dump gradient statistic data.
         level0: [step, max, min, norm, shape_dim, shape]
         level1: [step, max, min, norm, shape_dim, shape] + grad_bool_data
         level2: [step, max, min, norm, shape_dim, shape, dist_dim, dist] + grad_bool_data
-    '''
+    """
     dump_path = os.path.join(dump_dir, g_name)
     dump_dir_path = dump_path + "_dir"
     save_op = ms.ops.TensorDump()
@@ -182,7 +195,7 @@ class CSVGenerator(Process):
         shape_dim = int(stat_data[GradConst.SHAPE_DIM_IDX])
         file_name = os.path.basename(file_path)
         prefix_idx = len(file_name.split("_")[0])
-        param_name = file_name[(prefix_idx + 1) : -(len(GradConst.NPY_SUFFIX) + 1)]
+        param_name = file_name[(prefix_idx + 1): -(len(GradConst.NPY_SUFFIX) + 1)]
         if not param_name:
             raise RuntimeError("Invalid gradient statistic file name.")
         csv_line = [param_name]
@@ -224,8 +237,9 @@ class CSVGenerator(Process):
             if i == 0:
                 intervals.append(f"(-inf, {self.bounds[i]}]")
             else:
-                intervals.append(f"({self.bounds[i-1]}, {self.bounds[i]}]")
+                intervals.append(f"({self.bounds[i - 1]}, {self.bounds[i]}]")
         intervals.extend([f"({self.bounds[-1]}, inf)", "=0"])
         return intervals
 
+
 csv_generator = CSVGenerator()

msprobe/mindspore/grad_probe/grad_monitor.py

@@ -1,7 +1,22 @@
+# Copyright (c) 2024-2024, Huawei Technologies Co., Ltd.
+# All rights reserved.
+#
+# 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.
+
+from msprobe.core.grad_probe.constant import GradConst
 from msprobe.mindspore.grad_probe.global_context import grad_context
 from msprobe.mindspore.grad_probe.grad_analyzer import csv_generator
 from msprobe.mindspore.grad_probe.hook import hook_optimizer
-from msprobe.core.grad_probe.constant import GradConst
 
 
 class GradientMonitor: