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

msprobe/pytorch/compare/match.py

@@ -1,33 +1,34 @@
-import os
-from msprobe.core.common.utils import CompareException, load_yaml
-
-
-class AtenIrMapping():
-    def __init__(self):
-        cur_path = os.path.dirname(os.path.realpath(__file__))
-        yaml_path = os.path.join(cur_path, "mapping.yaml")
-        self.aten_mapping = load_yaml(yaml_path)
-    
-    def match(self, op1, op2):
-        if "Aten" in op1 and "Aten" not in op2:
-            return self.match_op(op1, op2)
-        else:
-            return self.match_op(op2, op1)
-
-    def match_op(self, aten_op, torch_op):
-        try:
-            aten_op_raw_name_overload = '_'.join(aten_op.split("_")[1:-3])
-            aten_op_raw_name = aten_op_raw_name_overload.split('.')[0]
-            torch_op_raw_name = '_'.join(torch_op.split("_")[1:-3]).lower()
-        except IndexError as e:
-            err_msg = f"Dump op name format error: {aten_op}, {torch_op}. Your dump data may be corrupted."
-            raise CompareException.INVALID_DATA_ERROR(err_msg) from e
-        matching_op = self.aten_mapping.get(aten_op_raw_name)
-        if matching_op is None:
-            return False
-        if matching_op.lower() == torch_op_raw_name:
-            return True
-        return False
-
-
-graph_mapping = AtenIrMapping()
+import os
+from msprobe.core.common.utils import CompareException
+from msprobe.core.common.file_utils import load_yaml
+
+
+class AtenIrMapping():
+    def __init__(self):
+        cur_path = os.path.dirname(os.path.realpath(__file__))
+        yaml_path = os.path.join(cur_path, "mapping.yaml")
+        self.aten_mapping = load_yaml(yaml_path)
+    
+    def match(self, op1, op2):
+        if "Aten" in op1 and "Aten" not in op2:
+            return self.match_op(op1, op2)
+        else:
+            return self.match_op(op2, op1)
+
+    def match_op(self, aten_op, torch_op):
+        try:
+            aten_op_raw_name_overload = '_'.join(aten_op.split("_")[1:-3])
+            aten_op_raw_name = aten_op_raw_name_overload.split('.')[0]
+            torch_op_raw_name = '_'.join(torch_op.split("_")[1:-3]).lower()
+        except IndexError as e:
+            err_msg = f"Dump op name format error: {aten_op}, {torch_op}. Your dump data may be corrupted."
+            raise CompareException.INVALID_DATA_ERROR(err_msg) from e
+        matching_op = self.aten_mapping.get(aten_op_raw_name)
+        if matching_op is None:
+            return False
+        if matching_op.lower() == torch_op_raw_name:
+            return True
+        return False
+
+
+graph_mapping = AtenIrMapping()

msprobe/pytorch/compare/pt_compare.py

@@ -1,40 +1,50 @@
-import os.path
-import torch
-from msprobe.core.common.const import FileCheckConst, Const
-from msprobe.core.common.log import logger
-from msprobe.core.common.exceptions import FileCheckException
-from msprobe.core.compare.acc_compare import Comparator 
-from msprobe.core.common.utils import create_directory, check_configuration_param, task_dumppath_get, \
-    check_compare_param, FileChecker
-from msprobe.core.common.utils import CompareException
-
-
-class PTComparator (Comparator):
-    def __init__(self):
-        self.frame_name = PTComparator.__name__
-    
-    def read_npy_data(self, dir_path, file_name):
-        data_path = os.path.join(dir_path, file_name)
-        path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE,
-                                FileCheckConst.PT_SUFFIX, False)
-        data_path = path_checker.common_check()
-        data_value = torch.load(data_path, map_location=torch.device('cpu')).detach()       # detach for less memory
-        if data_value.dtype == torch.bfloat16:
-            data_value = data_value.to(torch.float32)
-        data_value = data_value.numpy()
-        return data_value  
-    
-    
-def compare(input_param, output_path, stack_mode=False, auto_analyze=True, fuzzy_match=False):
-    try:
-        summary_compare, md5_compare = task_dumppath_get(input_param)
-        check_configuration_param(stack_mode, auto_analyze, fuzzy_match)
-        create_directory(output_path)
-        check_compare_param(input_param, output_path, summary_compare, md5_compare)
-    except (CompareException, FileCheckException) as error:
-        logger.error('Compare failed. Please check the arguments and do it again!')
-        raise CompareException(error.code) from error
-    pt_comparator = PTComparator()
-    pt_comparator.compare_core(input_param, output_path, stack_mode=stack_mode,
-                 auto_analyze=auto_analyze, fuzzy_match=fuzzy_match, summary_compare=summary_compare,
-                 md5_compare=md5_compare)
+import os.path
+import torch
+from msprobe.core.common.const import FileCheckConst
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.exceptions import FileCheckException
+from msprobe.core.compare.acc_compare import Comparator
+from msprobe.core.common.utils import check_configuration_param, task_dumppath_get, check_compare_param, CompareException
+from msprobe.core.common.file_utils import FileChecker, create_directory
+from msprobe.pytorch.common.utils import load_pt
+
+
+class PTComparator (Comparator):
+    def __init__(self):
+        self.frame_name = PTComparator.__name__
+    
+    def read_npy_data(self, dir_path, file_name):
+        data_path = os.path.join(dir_path, file_name)
+        path_checker = FileChecker(data_path, FileCheckConst.FILE, FileCheckConst.READ_ABLE,
+                                FileCheckConst.PT_SUFFIX, False)
+        data_path = path_checker.common_check()
+        try:
+            data_value = load_pt(data_path,
+                                 to_cpu=True).detach()  # detach because numpy can not process gradient information
+        except RuntimeError as e:
+            # 这里捕获 load_pt 中抛出的异常
+            logger.error(f"Failed to load the .pt file at {data_path}.")
+            raise CompareException(CompareException.INVALID_FILE_ERROR) from e
+        except AttributeError as e:
+            # 这里捕获 detach 方法抛出的异常
+            logger.error(f"Failed to detach the loaded tensor.")
+            raise CompareException(CompareException.DETACH_ERROR) from e
+        if data_value.dtype == torch.bfloat16:
+            data_value = data_value.to(torch.float32)
+        data_value = data_value.numpy()
+        return data_value  
+    
+    
+def compare(input_param, output_path, stack_mode=False, auto_analyze=True, fuzzy_match=False):
+    try:
+        summary_compare, md5_compare = task_dumppath_get(input_param)
+        check_configuration_param(stack_mode, auto_analyze, fuzzy_match)
+        create_directory(output_path)
+        check_compare_param(input_param, output_path, summary_compare, md5_compare)
+    except (CompareException, FileCheckException) as error:
+        logger.error('Compare failed. Please check the arguments and do it again!')
+        raise CompareException(error.code) from error
+    pt_comparator = PTComparator()
+    pt_comparator.compare_core(input_param, output_path, stack_mode=stack_mode,
+                 auto_analyze=auto_analyze, fuzzy_match=fuzzy_match, summary_compare=summary_compare,
+                 md5_compare=md5_compare)

msprobe/pytorch/debugger/debugger_config.py

@@ -1,95 +1,95 @@
-from msprobe.pytorch.common import seed_all
-from msprobe.pytorch.common.log import logger
-from msprobe.core.common.const import Const
-
-
-class DebuggerConfig:
-    def __init__(self, common_config, task_config, task, dump_path, level):
-        self.dump_path = dump_path if dump_path else common_config.dump_path
-        self.task = task or common_config.task or Const.STATISTICS
-        self.rank = common_config.rank if common_config.rank else []
-        self.step = common_config.step if common_config.step else []
-        self.level = level or common_config.level or "L1"
-        self.seed = common_config.seed if common_config.seed else 1234
-        self.is_deterministic = common_config.is_deterministic
-        self.enable_dataloader = common_config.enable_dataloader
-        self.scope = task_config.scope if task_config.scope else []
-        self.list = task_config.list if task_config.list else []
-        self.data_mode = task_config.data_mode if task_config.data_mode else ["all"]
-        self.backward_input_list = task_config.backward_input if task_config.backward_input else []
-        self.backward_input = {}
-        self.acl_config = common_config.acl_config if common_config.acl_config else ""
-        self.is_forward_acl_dump = True
-        self.summary_mode = task_config.summary_mode if task_config.summary_mode else Const.STATISTICS
-        self.overflow_nums = task_config.overflow_nums if task_config.overflow_nums else 1
-        self.framework = Const.PT_FRAMEWORK
-
-        if self.task == Const.FREE_BENCHMARK:
-            self.fuzz_device = task_config.fuzz_device if task_config.fuzz_device else 'npu'
-            self.handler_type = task_config.handler_type if task_config.handler_type else 'check'
-            self.pert_mode = task_config.pert_mode if task_config.pert_mode else 'improve_precision'
-            self.fuzz_level = task_config.fuzz_level if task_config.fuzz_level else 'L1'
-            self.fuzz_stage = task_config.fuzz_stage if task_config.fuzz_stage else 'forward'
-            self.preheat_config = {
-                "if_preheat": task_config.if_preheat if task_config.if_preheat is not None else True, 
-                "preheat_step": task_config.preheat_step if task_config.preheat_step else 15, 
-                "max_sample": task_config.max_sample if task_config.max_sample else 20, 
-            }
-
-        self.online_run_ut = False
-        if self.task == Const.TENSOR:
-            # dump api tensor and collaborate with online run_ut
-            self.online_run_ut = task_config.online_run_ut if task_config.online_run_ut else False
-            self.nfs_path = task_config.nfs_path if task_config.nfs_path else ""
-            self.tls_path = task_config.tls_path if task_config.tls_path else ""
-            self.host = task_config.host if task_config.host else ""
-            self.port = task_config.port if task_config.port else -1
-
-        self.check()
-        if self.step:
-            self.step.sort()
-        if self.level == "L2":
-            if not self.scope or not isinstance(self.scope, list) or len(self.scope) != 1:
-                raise ValueError("scope must be configured as a list with one api name")
-            if isinstance(self.scope[0], str) and Const.BACKWARD in self.scope[0] and not self.backward_input_list:
-                raise ValueError("backward_input must be configured when scope contains 'backward'")
-            if Const.BACKWARD in self.scope[0]:
-                self.is_forward_acl_dump = False
-                for index, scope_spec in enumerate(self.scope):
-                    self.scope[index] = scope_spec.replace(Const.BACKWARD, Const.FORWARD)
-                    self.backward_input[self.scope[index]] = self.backward_input_list[index]
-        seed_all(self.seed, self.is_deterministic)
-
-    def check_kwargs(self):
-        if self.task and self.task not in Const.TASK_LIST:
-            raise Exception("task is invalid")
-        if self.level and self.level not in Const.LEVEL_LIST:
-            raise Exception("level is invalid")
-        if not self.dump_path:
-            raise Exception("Invalid dump path, please check your config")
-
-    def check(self):
-        self.check_kwargs()
-        self._check_rank()
-        self._check_step()
-        return True
-
-    def check_model(self, model):
-        if self.level in ["L0", "mix"] and not model:
-            raise Exception(
-                f"For level {self.level}, PrecisionDebugger must receive a model argument."
-            )
-
-    def _check_rank(self):
-        if self.rank:
-            for rank_id in self.rank:
-                if not isinstance(rank_id, int) or rank_id < 0:
-                    raise ValueError(f"rank {self.rank} must be an integer and greater than or equal to 0.")
-            else:
-                logger.warning_on_rank_0(f"Rank argument is provided. Only rank {self.rank} data will be dumpped.")
-
-    def _check_step(self):
-        if self.step:
-            for s in self.step:
-                if not isinstance(s, int) or s < 0:
-                    raise ValueError(f"step element {s} must be an integer and greater than or equal to 0.")
+from msprobe.pytorch.common import seed_all
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.const import Const
+
+
+class DebuggerConfig:
+    def __init__(self, common_config, task_config, task, dump_path, level):
+        self.dump_path = dump_path if dump_path else common_config.dump_path
+        self.task = task or common_config.task or Const.STATISTICS
+        self.rank = common_config.rank if common_config.rank else []
+        self.step = common_config.step if common_config.step else []
+        self.level = level or common_config.level or "L1"
+        self.seed = common_config.seed if common_config.seed else 1234
+        self.is_deterministic = common_config.is_deterministic
+        self.enable_dataloader = common_config.enable_dataloader
+        self.scope = task_config.scope if task_config.scope else []
+        self.list = task_config.list if task_config.list else []
+        self.data_mode = task_config.data_mode if task_config.data_mode else ["all"]
+        self.backward_input_list = task_config.backward_input if task_config.backward_input else []
+        self.backward_input = {}
+        self.acl_config = common_config.acl_config if common_config.acl_config else ""
+        self.is_forward_acl_dump = True
+        self.summary_mode = task_config.summary_mode if task_config.summary_mode else Const.STATISTICS
+        self.overflow_nums = task_config.overflow_nums if task_config.overflow_nums else 1
+        self.framework = Const.PT_FRAMEWORK
+
+        if self.task == Const.FREE_BENCHMARK:
+            self.fuzz_device = task_config.fuzz_device if task_config.fuzz_device else 'npu'
+            self.handler_type = task_config.handler_type if task_config.handler_type else 'check'
+            self.pert_mode = task_config.pert_mode if task_config.pert_mode else 'improve_precision'
+            self.fuzz_level = task_config.fuzz_level if task_config.fuzz_level else 'L1'
+            self.fuzz_stage = task_config.fuzz_stage if task_config.fuzz_stage else 'forward'
+            self.preheat_config = {
+                "if_preheat": task_config.if_preheat if task_config.if_preheat is not None else True, 
+                "preheat_step": task_config.preheat_step if task_config.preheat_step else 15, 
+                "max_sample": task_config.max_sample if task_config.max_sample else 20, 
+            }
+
+        self.online_run_ut = False
+        if self.task == Const.TENSOR:
+            # dump api tensor and collaborate with online run_ut
+            self.online_run_ut = task_config.online_run_ut if task_config.online_run_ut else False
+            self.nfs_path = task_config.nfs_path if task_config.nfs_path else ""
+            self.tls_path = task_config.tls_path if task_config.tls_path else ""
+            self.host = task_config.host if task_config.host else ""
+            self.port = task_config.port if task_config.port else -1
+
+        self.check()
+        if self.step:
+            self.step.sort()
+        if self.level == "L2":
+            if not self.scope or not isinstance(self.scope, list) or len(self.scope) != 1:
+                raise ValueError("scope must be configured as a list with one api name")
+            if isinstance(self.scope[0], str) and Const.BACKWARD in self.scope[0] and not self.backward_input_list:
+                raise ValueError("backward_input must be configured when scope contains 'backward'")
+            if Const.BACKWARD in self.scope[0]:
+                self.is_forward_acl_dump = False
+                for index, scope_spec in enumerate(self.scope):
+                    self.scope[index] = scope_spec.replace(Const.BACKWARD, Const.FORWARD)
+                    self.backward_input[self.scope[index]] = self.backward_input_list[index]
+        seed_all(self.seed, self.is_deterministic)
+
+    def check_kwargs(self):
+        if self.task and self.task not in Const.TASK_LIST:
+            raise Exception("task is invalid")
+        if self.level and self.level not in Const.LEVEL_LIST:
+            raise Exception("level is invalid")
+        if not self.dump_path:
+            raise Exception("Invalid dump path, please check your config")
+
+    def check(self):
+        self.check_kwargs()
+        self._check_rank()
+        self._check_step()
+        return True
+
+    def check_model(self, model):
+        if self.level in ["L0", "mix"] and not model:
+            raise Exception(
+                f"For level {self.level}, PrecisionDebugger must receive a model argument."
+            )
+
+    def _check_rank(self):
+        if self.rank:
+            for rank_id in self.rank:
+                if not isinstance(rank_id, int) or rank_id < 0:
+                    raise ValueError(f"rank {self.rank} must be an integer and greater than or equal to 0.")
+            else:
+                logger.warning_on_rank_0(f"Rank argument is provided. Only rank {self.rank} data will be dumpped.")
+
+    def _check_step(self):
+        if self.step:
+            for s in self.step:
+                if not isinstance(s, int) or s < 0:
+                    raise ValueError(f"step element {s} must be an integer and greater than or equal to 0.")

msprobe/pytorch/debugger/precision_debugger.py

@@ -1,125 +1,125 @@
-import torch
-from torch.utils.data import dataloader
-from msprobe.pytorch.debugger.debugger_config import DebuggerConfig
-from msprobe.pytorch.service import Service
-from msprobe.pytorch.common.log import logger
-from msprobe.pytorch.pt_config import parse_json_config
-from msprobe.core.common.exceptions import MsprobeException
-from msprobe.core.common.const import Const
-from msprobe.pytorch.grad_probe.grad_monitor import GradientMonitor
-
-
-class PrecisionDebugger:
-    _instance = None
-    tasks_not_need_debugger = [Const.GRAD_PROBE]
-
-    def __new__(cls, *args, **kwargs):
-        if cls._instance is None:
-            cls._instance = super(PrecisionDebugger, cls).__new__(cls)
-            cls._instance.config = None
-            cls._instance.enable_dataloader = False
-        return cls._instance
-
-    def __init__(
-        self,
-        config_path=None,
-        task=None,
-        dump_path=None,
-        level=None,
-        model=None,
-        step=None,
-    ):
-        if not hasattr(self, "initialized"):
-            self.api_origin = False
-            self.initialized = True
-            self.model = self.check_model_valid(model)
-            common_config, task_config = parse_json_config(config_path, task)
-            self.task = common_config.task
-            if self.task == Const.GRAD_PROBE:
-                self.gm = GradientMonitor(common_config, task_config)
-                return
-            if step:
-                common_config.step = step
-            self.config = DebuggerConfig(
-                common_config, task_config, task, dump_path, level
-            )
-            self.config.check_model(self.model)
-            self.service = Service(self.config)
-            self.enable_dataloader = self.config.enable_dataloader
-            if self.enable_dataloader:
-                logger.warning_on_rank_0("The enable_dataloader feature will be deprecated in the future.")
-                dataloader._BaseDataLoaderIter.__next__ = iter_tracer(dataloader._BaseDataLoaderIter.__next__)
-
-    @property
-    def instance(self):
-        return self._instance
-
-    @staticmethod
-    def check_model_valid(model):
-        if not model or isinstance(model, torch.nn.Module):
-            return model
-        raise MsprobeException(
-            MsprobeException.INVALID_PARAM_ERROR, "model 参数必须是torch.nn.Module类型。"
-        )
-
-    @classmethod
-    def start(cls):
-        instance = cls._instance
-        if instance.task in PrecisionDebugger.tasks_not_need_debugger:
-            return
-        if not instance:
-            raise Exception("No instance of PrecisionDebugger found.")
-        if instance.enable_dataloader:
-            logger.warning_on_rank_0("DataLoader is enabled, start() skipped.")
-        else:
-            instance.service.start(instance.model, instance.api_origin)
-            instance.api_origin = False
-
-    # 指定代码段dump前反向结束符，之后的计算过程数据将被忽略，无法被dump
-    @classmethod
-    def forward_backward_dump_end(cls):
-        instance = cls._instance
-        instance.service.forward_backward_dump_end()
-        instance.api_origin = True
-
-    @classmethod
-    def stop(cls):
-        instance = cls._instance
-        if instance.task in PrecisionDebugger.tasks_not_need_debugger:
-            return
-        if not instance:
-            raise Exception("PrecisionDebugger instance is not created.")
-        if instance.enable_dataloader:
-            logger.warning_on_rank_0("DataLoader is enabled, stop() skipped.")
-        else:
-            instance.service.stop()
-
-    @classmethod
-    def step(cls):
-        if cls._instance.task in PrecisionDebugger.tasks_not_need_debugger:
-            return
-        if not cls._instance:
-            raise Exception("PrecisionDebugger instance is not created.")
-        cls._instance.service.step()
-
-    @classmethod
-    def monitor(cls, model):
-        if not cls._instance:
-            raise Exception("PrecisionDebugger instance is not created.")
-        if cls._instance.task != Const.GRAD_PROBE:
-            return
-        cls._instance.gm.monitor(model)
-
-
-def iter_tracer(func):
-    def func_wrapper(*args, **kwargs):
-        debugger_instance = PrecisionDebugger.instance
-        debugger_instance.enable_dataloader = False
-        if not debugger_instance.service.first_start:
-            debugger_instance.stop()
-            debugger_instance.step()
-        result = func(*args, **kwargs)
-        debugger_instance.start()
-        debugger_instance.enable_dataloader = True
-        return result
-    return func_wrapper
+import torch
+from torch.utils.data import dataloader
+from msprobe.pytorch.debugger.debugger_config import DebuggerConfig
+from msprobe.pytorch.service import Service
+from msprobe.pytorch.common.log import logger
+from msprobe.pytorch.pt_config import parse_json_config
+from msprobe.core.common.exceptions import MsprobeException
+from msprobe.core.common.const import Const
+from msprobe.pytorch.grad_probe.grad_monitor import GradientMonitor
+
+
+class PrecisionDebugger:
+    _instance = None
+    tasks_not_need_debugger = [Const.GRAD_PROBE]
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super(PrecisionDebugger, cls).__new__(cls)
+            cls._instance.config = None
+            cls._instance.enable_dataloader = False
+        return cls._instance
+
+    def __init__(
+        self,
+        config_path=None,
+        task=None,
+        dump_path=None,
+        level=None,
+        model=None,
+        step=None,
+    ):
+        if not hasattr(self, "initialized"):
+            self.api_origin = False
+            self.initialized = True
+            self.model = self.check_model_valid(model)
+            common_config, task_config = parse_json_config(config_path, task)
+            self.task = common_config.task
+            if self.task == Const.GRAD_PROBE:
+                self.gm = GradientMonitor(common_config, task_config)
+                return
+            if step:
+                common_config.step = step
+            self.config = DebuggerConfig(
+                common_config, task_config, task, dump_path, level
+            )
+            self.config.check_model(self.model)
+            self.service = Service(self.config)
+            self.enable_dataloader = self.config.enable_dataloader
+            if self.enable_dataloader:
+                logger.warning_on_rank_0("The enable_dataloader feature will be deprecated in the future.")
+                dataloader._BaseDataLoaderIter.__next__ = iter_tracer(dataloader._BaseDataLoaderIter.__next__)
+
+    @property
+    def instance(self):
+        return self._instance
+
+    @staticmethod
+    def check_model_valid(model):
+        if not model or isinstance(model, torch.nn.Module):
+            return model
+        raise MsprobeException(
+            MsprobeException.INVALID_PARAM_ERROR, "model 参数必须是torch.nn.Module类型。"
+        )
+
+    @classmethod
+    def start(cls):
+        instance = cls._instance
+        if instance.task in PrecisionDebugger.tasks_not_need_debugger:
+            return
+        if not instance:
+            raise Exception("No instance of PrecisionDebugger found.")
+        if instance.enable_dataloader:
+            logger.warning_on_rank_0("DataLoader is enabled, start() skipped.")
+        else:
+            instance.service.start(instance.model, instance.api_origin)
+            instance.api_origin = False
+
+    # 指定代码段dump前反向结束符，之后的计算过程数据将被忽略，无法被dump
+    @classmethod
+    def forward_backward_dump_end(cls):
+        instance = cls._instance
+        instance.service.forward_backward_dump_end()
+        instance.api_origin = True
+
+    @classmethod
+    def stop(cls):
+        instance = cls._instance
+        if instance.task in PrecisionDebugger.tasks_not_need_debugger:
+            return
+        if not instance:
+            raise Exception("PrecisionDebugger instance is not created.")
+        if instance.enable_dataloader:
+            logger.warning_on_rank_0("DataLoader is enabled, stop() skipped.")
+        else:
+            instance.service.stop()
+
+    @classmethod
+    def step(cls):
+        if cls._instance.task in PrecisionDebugger.tasks_not_need_debugger:
+            return
+        if not cls._instance:
+            raise Exception("PrecisionDebugger instance is not created.")
+        cls._instance.service.step()
+
+    @classmethod
+    def monitor(cls, model):
+        if not cls._instance:
+            raise Exception("PrecisionDebugger instance is not created.")
+        if cls._instance.task != Const.GRAD_PROBE:
+            return
+        cls._instance.gm.monitor(model)
+
+
+def iter_tracer(func):
+    def func_wrapper(*args, **kwargs):
+        debugger_instance = PrecisionDebugger.instance
+        debugger_instance.enable_dataloader = False
+        if not debugger_instance.service.first_start:
+            debugger_instance.stop()
+            debugger_instance.step()
+        result = func(*args, **kwargs)
+        debugger_instance.start()
+        debugger_instance.enable_dataloader = True
+        return result
+    return func_wrapper

msprobe/pytorch/free_benchmark/__init__.py

@@ -1,8 +1,8 @@
-from msprobe.core.common.log import logger
-from msprobe.core.common.exceptions import FreeBenchmarkException
-from msprobe.core.common.const import Const
-
-from .main import FreeBenchmarkCheck
-from .common.params import UnequalRow
-
-__all__ = [FreeBenchmarkCheck, UnequalRow]
+from msprobe.pytorch.common.log import logger
+from msprobe.core.common.exceptions import FreeBenchmarkException
+from msprobe.core.common.const import Const
+
+from .main import FreeBenchmarkCheck
+from .common.params import UnequalRow
+
+__all__ = [FreeBenchmarkCheck, UnequalRow]