mindstudio-probe 1.0.1__py3-none-any.whl

msprobe/test/pytorch_ut/api_accuracy_checker/run_ut/test_run_ut.py

@@ -0,0 +1,72 @@
+# coding=utf-8
+import os
+import copy
+import unittest
+import torch
+from unittest.mock import patch, DEFAULT
+from msprobe.pytorch.api_accuracy_checker.run_ut.run_ut import *
+from msprobe.pytorch.api_accuracy_checker.common.utils import get_json_contents
+
+base_dir = os.path.dirname(os.path.realpath(__file__))
+forward_file = os.path.join(base_dir, "forward.json")
+forward_content = get_json_contents(forward_file)
+for api_full_name, api_info_dict in forward_content.items():
+    api_full_name = api_full_name
+    api_info_dict = api_info_dict
+
+
+class TestRunUtMethods(unittest.TestCase):
+    def test_exec_api(self):
+        api_info = copy.deepcopy(api_info_dict)
+
+        [api_type, api_name, _, _] = api_full_name.split(".")
+        args, kwargs, need_grad = get_api_info(api_info, api_name, None)
+        cpu_args, cpu_kwargs = generate_cpu_params(args, kwargs, True, '')
+        out = exec_api(api_type, api_name, cpu_args, cpu_kwargs)
+        self.assertEqual(out[0].dtype, torch.float32)
+        self.assertTrue(out[0].requires_grad)
+        self.assertEqual(out[0].shape, torch.Size([2048, 2, 1, 128]))
+
+    def test_generate_device_params(self):
+        mock_tensor = torch.rand([2, 2560, 24, 24], dtype=torch.float32, requires_grad=True)
+
+        with patch.multiple('torch.Tensor',
+                            to=DEFAULT,
+                            clone=DEFAULT,
+                            detach=DEFAULT,
+                            requires_grad_=DEFAULT,
+                            type_as=DEFAULT,
+                            retain_grad=DEFAULT) as mocks:
+            mocks['clone'].return_value = mock_tensor
+            mocks['detach'].return_value = mock_tensor
+            mocks['requires_grad_'].return_value = mock_tensor
+            mocks['type_as'].return_value = mock_tensor
+            mocks['retain_grad'].return_value = None
+            mocks['to'].return_value = mock_tensor
+
+            device_args, device_kwargs = generate_device_params([mock_tensor], {'inplace': False}, True, '')
+            self.assertEqual(len(device_args), 1)
+            self.assertEqual(device_args[0].dtype, torch.float32)
+            self.assertTrue(device_args[0].requires_grad)
+            self.assertEqual(device_args[0].shape, torch.Size([2, 2560, 24, 24]))
+            self.assertEqual(device_kwargs, {'inplace': False})
+
+    def test_generate_cpu_params(self):
+        api_info = copy.deepcopy(api_info_dict)
+        [api_type, api_name, _, _] = api_full_name.split(".")
+        args, kwargs, need_grad = get_api_info(api_info, api_name, None)
+        cpu_args, cpu_kwargs = generate_cpu_params(args, kwargs, True, '')
+        self.assertEqual(len(cpu_args), 2)
+        self.assertEqual(cpu_args[0].dtype, torch.float32)
+        self.assertTrue(cpu_args[0].requires_grad)
+        self.assertEqual(cpu_args[0].shape, torch.Size([2048, 2, 1, 256]))
+        self.assertEqual(cpu_kwargs, {'dim': -1})
+
+    def test_UtDataInfo(self):
+        data_info = UtDataInfo(None, None, None, None, None, None, None)
+        self.assertIsNone(data_info.bench_grad)
+        self.assertIsNone(data_info.device_grad)
+        self.assertIsNone(data_info.device_output)
+        self.assertIsNone(data_info.bench_output)
+        self.assertIsNone(data_info.grad_in)
+        self.assertIsNone(data_info.in_fwd_data_list)

msprobe/test/pytorch_ut/compare/test_acc_compare.py

@@ -0,0 +1,17 @@
+# coding=utf-8
+import unittest
+from msprobe.pytorch.compare.acc_compare import rename_api
+
+class TestUtilsMethods(unittest.TestCase):
+
+    def test_rename_api(self):
+        test_name_1 = "Distributed.broadcast.0.forward.input.0"
+        expect_name_1 = "Distributed.broadcast.input.0"
+        actual_name_1 = rename_api(test_name_1, "forward")
+        self.assertEqual(actual_name_1, expect_name_1)
+        
+        test_name_2 = "Torch.sum.0.backward.output.0"
+        expect_name_2 = "Torch.sum.output.0"
+        actual_name_2 = rename_api(test_name_2, "backward")
+        self.assertEqual(actual_name_2, expect_name_2)
+        

msprobe/test/pytorch_ut/free_benchmark/perturbed_layers/test_perturbed_layser.py

@@ -0,0 +1,105 @@
+from unittest import TestCase
+
+import torch
+from msprobe.core.common.const import Const
+from msprobe.pytorch.free_benchmark.common.enums import DeviceType, PerturbationMode
+from msprobe.pytorch.free_benchmark.common.params import data_pre_deal
+from msprobe.pytorch.free_benchmark.perturbed_layers.layer_factory import LayerFactory
+
+
+class TestPerturbedLayer(TestCase):
+
+    # 对输出精度和输入精度一致算子使用升精度扰动因子时, 输出结果的精度也会提升
+    def test_improve_precision_layer_handle_with_out_dtype_changing(self):
+        api_name = "Torch.mul.0.forward"
+        x = torch.randn(2, 3, dtype=torch.float16)
+        y = torch.randn(2, 3, dtype=torch.float16)
+        out = torch.mul(x, y)
+
+        data_params = data_pre_deal(api_name, torch.mul, (x, y), {})
+        data_params.fuzz_stage = Const.FORWARD
+        data_params.original_result = out
+
+        layer = LayerFactory.create(
+            api_name, DeviceType.NPU, PerturbationMode.IMPROVE_PRECISION
+        )
+        layer.handle(data_params)
+        self.assertEqual(data_params.original_result.dtype, torch.float16)
+        self.assertEqual(layer.perturbed_value, torch.float32)
+        self.assertEqual(data_params.perturbed_result.dtype, torch.float32)
+
+    # 对于可迭代类型的输入, 升精度方法会遍历其中元素对支持类型输入升精度
+    def test_improve_precision_layer_with_iterable_inputs(self):
+        api_name = "iterable.0.forward"
+        tensor_a = torch.randn(2, 3, dtype=torch.bfloat16)
+        tensor_b = torch.randn(2, 3, dtype=torch.float16)
+        tensor_c = torch.randn(2, 3, dtype=torch.float32)
+        tensor_d = torch.randn(2, 3, dtype=torch.float64)
+        tensor_f = torch.randn(2, 3, dtype=torch.float64).to(torch.int32)
+        inputs = [tensor_a, tensor_b, {"c": tensor_c, "d": tensor_d}, tensor_f]
+
+        layer = LayerFactory.create(
+            api_name, DeviceType.NPU, PerturbationMode.IMPROVE_PRECISION
+        )
+        Perturbed_value = layer.improve_tensor_precision(inputs)
+        self.assertEqual(Perturbed_value[0].dtype, torch.float32)
+        self.assertEqual(Perturbed_value[1].dtype, torch.float32)
+        self.assertEqual(Perturbed_value[2]["c"].dtype, torch.float32)
+        self.assertEqual(Perturbed_value[2]["d"].dtype, torch.float64)
+        self.assertEqual(Perturbed_value[3].dtype, torch.int32)
+
+    # no_change扰动因子不会改变输入
+    def test_no_change_layer(self):
+        api_name = "nochange.0.forward"
+        inputs = torch.as_tensor([1e-9, 1e-2], dtype=torch.float32)
+        layer = LayerFactory.create(
+            api_name, DeviceType.NPU, PerturbationMode.NO_CHANGE
+        )
+        Perturbed_value = layer.no_change(inputs)
+        self.assertEqual(Perturbed_value[0], 1e-9)
+        self.assertEqual(Perturbed_value[1], 1e-2)
+
+    # 对于一维二维张量，change_value扰动因子会交换首尾值的位置
+    def test_change_value_layer(self):
+        api_name = "change.0.forward"
+        inputs_1dim = torch.as_tensor([1e-9, 1e-7, 1e-2], dtype=torch.float32)
+        inputs_2dim = torch.as_tensor(
+            [[1e-9, 1e-7, 1e-2], [1e-9, 1e-2, 1e-7]], dtype=torch.float32
+        )
+        layer = LayerFactory.create(
+            api_name, DeviceType.NPU, PerturbationMode.CHANGE_VALUE
+        )
+        Perturbed_value_1dim = layer.change_value(inputs_1dim)
+        layer.is_added = False
+        Perturbed_value_2dim = layer.change_value(inputs_2dim)
+        self.assertEqual(Perturbed_value_1dim[0], 1e-2)
+        self.assertEqual(Perturbed_value_1dim[2], 1e-9)
+        self.assertEqual(Perturbed_value_2dim[0][0], 1e-7)
+        self.assertEqual(Perturbed_value_2dim[-1][-1], 1e-9)
+
+    # 对于输入张量，bit_noise扰动因子对大于极小值的部分进行末尾比特翻转
+    def test_bit_noise_layer(self):
+        api_name = "bitnoise.0.forward"
+        inputs = torch.as_tensor(
+            [4096.00048828125, 16777216, 1e-38], dtype=torch.float32
+        )
+        layer = LayerFactory.create(
+            api_name, DeviceType.NPU, PerturbationMode.BIT_NOISE
+        )
+        Perturbed_value = layer.add_bit_noise(inputs)
+        self.assertEqual(Perturbed_value[0], 4096.0000000000)
+        self.assertEqual(Perturbed_value[1], 16777218)
+        self.assertEqual(Perturbed_value[2], 1e-38)
+
+    # 对于输入张量，add_noise扰动因子对大于极小值的部分增加一个小值
+    def test_add_noise_layer(self):
+        api_name = "addnoise.0.forward"
+        inputs = torch.as_tensor(
+            [1e-1, 1e-2], dtype=torch.bfloat16
+        )
+        layer = LayerFactory.create(
+            api_name, DeviceType.NPU, PerturbationMode.ADD_NOISE
+        )
+        Perturbed_value = layer.add_noise(inputs)
+        self.assertEqual(Perturbed_value[0], 1e-1+1e-4)
+        self.assertEqual(Perturbed_value[1], 1e-2)

msprobe/test/pytorch_ut/free_benchmark/result_handlers/test_result_handler.py

@@ -0,0 +1,121 @@
+from abc import ABC
+from unittest import TestCase
+
+import torch
+from msprobe.core.common.const import Const
+from msprobe.pytorch.free_benchmark.common.constant import PreheatConfig, ThresholdConfig
+from msprobe.pytorch.free_benchmark.common.counter import preheat_counter
+from msprobe.pytorch.free_benchmark.common.enums import (
+    DeviceType,
+    FuzzLevel,
+    HandlerType,
+    PerturbationMode,
+)
+from msprobe.pytorch.free_benchmark.common.params import DataParams, make_handler_params
+from msprobe.pytorch.free_benchmark.result_handlers.handler_factory import (
+    FuzzHandlerFactory,
+)
+
+
+class Config(ABC):
+    """
+    用以提供参数配置
+    """
+    def __init__(self, handler_type, preheat_config):
+        self.fuzz_stage = Const.FORWARD
+        self.handler_type = handler_type
+        self.fuzz_device = DeviceType.NPU
+        self.fuzz_level = FuzzLevel.BASE_LEVEL
+        self.pert_mode = PerturbationMode.IMPROVE_PRECISION
+        self.preheat_config = preheat_config
+
+
+class TestFuzzHandler(TestCase):
+
+    def setUp(self) -> None:
+        origin_inputs = [
+            torch.as_tensor([3.01, 3.02], dtype=torch.float16),
+            torch.as_tensor([0.02, 0.02], dtype=torch.float16),
+        ]
+        # 将输入乘以一个大于误差阈值1.002的值，模拟二次执行出现误差
+        perturbed_inputs = [
+            (value * 1.0021).to(torch.float32).to("cpu") for value in origin_inputs
+        ]
+        origin_output = torch.add(*origin_inputs)
+        perturbed_output = torch.add(*perturbed_inputs)
+        # 实例有问题的data对象
+        self.data_params = DataParams(
+            args=origin_inputs,
+            kwargs={},
+            original_result=origin_output,
+            perturbed_result=perturbed_output,
+            origin_func=torch.add,
+        )
+        self.api_name = "add.0.forward"
+        self.step = 0
+
+    def test_result_handler_check(self):
+        # 对于check处理类，扰动前后输出不一致的情况会有UnequalRow对象生成
+        for _ in range(2):
+            config = Config(
+                HandlerType.CHECK, {PreheatConfig.IF_PREHEAT: False}
+            )
+            handler_params = make_handler_params(self.api_name, config, self.step)
+            handler = FuzzHandlerFactory.create(handler_params)
+            handler.handle(self.data_params)
+            self.assertEqual(
+                len(handler.get_unequal_rows()), 1
+            )
+
+    def test_result_handler_fix(self):
+        # 对于fix处理类，扰动后输出会替代原始输出, dtype和原始输出一致，但值为新输出值
+        config = Config(
+            HandlerType.FIX, {PreheatConfig.IF_PREHEAT: False}
+        )
+        handler_params = make_handler_params(self.api_name, config, self.step)
+        handler = FuzzHandlerFactory.create(handler_params)
+        result = handler.handle(self.data_params)
+        self.assertEqual(result.dtype, torch.float16)
+        self.assertEqual(result.device, self.data_params.original_result.device)
+        self.assertAlmostEqual(
+            result[0], self.data_params.perturbed_result.to(torch.float16)[0]
+        )
+        self.assertAlmostEqual(
+            result[1], self.data_params.perturbed_result.to(torch.float16)[1]
+        )
+
+    def test_result_handler_preheat(self):
+        # 对于preheat处理类，在预热阶段后的阈值会根据CPU调整
+        config = Config(
+            HandlerType.CHECK,
+            {
+                PreheatConfig.IF_PREHEAT: True,
+                PreheatConfig.PREHEAT_STEP: 4,
+                PreheatConfig.MAX_SAMPLE: 3
+            }
+        )
+        for _ in range(3):
+            handler_params = make_handler_params(self.api_name, config, 0)
+            handler = FuzzHandlerFactory.create(handler_params)
+            handler.handle(self.data_params)
+        # 通过preheat_counter的数据可以判断预热是否正常执行，这里第一个step会记录api执行次数
+        self.assertEqual(preheat_counter.get_one_step_used_api("add"), 3)
+        for step in range(1, 4):
+            for _ in range(3):
+                handler_params = make_handler_params(self.api_name, config, step)
+                handler = FuzzHandlerFactory.create(handler_params)
+                handler.handle(self.data_params)
+            # call time记录当前step api的调用次数
+            self.assertEqual(preheat_counter.get_api_called_time("add"), 3)
+            # 对于3个step最多采样三次的预热设置，sample time应该每次采样一例
+            self.assertEqual(preheat_counter.get_api_sample_time("add"), 1)
+            # 预热阶段，api阈值应该在两个阈值超参之间
+            api_threshld = preheat_counter.get_api_thd("add", "torch.float16")
+            self.assertLessEqual(
+                api_threshld,
+                ThresholdConfig.PREHEAT_INITIAL_THD
+            )
+            self.assertGreaterEqual(
+                api_threshld,
+                ThresholdConfig.DTYPE_PER_THD[torch.float16]
+            )

msprobe/test/pytorch_ut/free_benchmark/test_main.py

@@ -0,0 +1,101 @@
+import functools
+from abc import ABC
+from unittest import TestCase
+
+import torch
+import torch.nn as nn
+from msprobe.core.common.const import Const
+from msprobe.pytorch.free_benchmark import FreeBenchmarkCheck
+from msprobe.pytorch.free_benchmark.common.constant import CommonField, PreheatConfig
+from msprobe.pytorch.free_benchmark.common.enums import (
+    DeviceType,
+    FuzzLevel,
+    HandlerType,
+    PerturbationMode,
+)
+
+
+class Config(ABC):
+    """
+    用以提供参数配置
+    """
+
+    def __init__(self, fuzz_stage, handler_type):
+        self.fuzz_stage = fuzz_stage
+        self.handler_type = handler_type
+        self.fuzz_device = DeviceType.NPU
+        self.fuzz_level = FuzzLevel.BASE_LEVEL
+        self.pert_mode = PerturbationMode.IMPROVE_PRECISION
+        self.preheat_config = {PreheatConfig.IF_PREHEAT: False}
+
+
+class WrapMul(nn.Module):
+    """
+    用nn.module包装mul算子, 在forward中调用torch.mul
+    """
+
+    def __init__(self, op_name) -> None:
+        super().__init__()
+        self.op_name = op_name
+
+    def forward(self, *args, **kwargs):
+        return torch.mul(*args, **kwargs)
+
+
+class UnequalDataProcessor(ABC):
+    """
+    接口类, 处理检测不一致结果
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.unequal_rows = []
+
+    def update_unequal_rows(self, unequal_rows):
+        self.unequal_rows.append(unequal_rows)
+
+
+class TestInterface(TestCase):
+    def setUp(self):
+        self.api_name = "Torch.mul.0"
+
+    def testForwardFix(self):
+        # 对于前向接口，在forward钩子中开启FIX，返回结果给hook的输出
+        config = Config(Const.FORWARD, HandlerType.FIX)
+        checker = FreeBenchmarkCheck(config)
+        # 执行算子前向
+        x = torch.randn(2, 3).to(torch.float16)
+        y = torch.randn(2, 3).to(torch.float16)
+        mul_module = WrapMul(self.api_name)
+        out = mul_module(x, y)
+        # 模拟forward hook中调用无标杆前向检测接口
+        result, _ = checker.forward(
+            self.api_name,
+            mul_module,
+            args=(x, y),
+            kwargs={},
+            output=out,
+        )
+        self.assertEqual(result.dtype, torch.float32)
+
+    def testBackwardCheck(self):
+        # 对于反向接口，在pre forward时暂存input, 然后在backwrad后进行对比
+        config = Config(Const.BACKWARD, HandlerType.CHECK)
+        checker = FreeBenchmarkCheck(config)
+        processor = UnequalDataProcessor()
+        # 初始化输入输出
+        x = torch.tensor([2, 3], dtype=torch.float16, requires_grad=True)
+        y = torch.tensor([2, 3], dtype=torch.float16, requires_grad=True)
+        grad_output =  torch.tensor([1,1], dtype=torch.float16)
+        backward_name = Const.SEP.join([self.api_name, Const.BACKWARD])
+        # 执行前向生成grad saver实例
+        mul_module = WrapMul(self.api_name)
+        checker.pre_forward(backward_name, mul_module, processor, (x, y), {})
+        # 执行算子前向和反向, 并反向获取扰动后grad_input
+        out = mul_module(x, y)
+        checker.backward(backward_name, mul_module, grad_output)
+        out.backward(torch.ones_like(out))
+        # module是否添加暂存器, 其中反向钩子执行扰动后grad_input是否正确
+        self.assertTrue(hasattr(mul_module, CommonField.GRADSAVER))
+        grad_saver = getattr(mul_module, CommonField.GRADSAVER)
+        self.assertEqual(grad_saver.perturbed_grad_input[0][0], 2)

msprobe/test/pytorch_ut/functional/test_dump_module.py

@@ -0,0 +1,15 @@
+import unittest
+
+import torch.nn as nn
+from msprobe.pytorch import PrecisionDebugger
+from msprobe.pytorch.functional.dump_module import module_dump, module_count
+
+
+class TestDumpModule(unittest.TestCase):
+    def setUp(self):
+        self.module = nn.Linear(in_features=8, out_features=4)
+
+    def test_module_dump(self):
+        PrecisionDebugger(dump_path="./dump")
+        module_dump(self.module, "TestModule")
+        self.assertTrue("TestModule" in module_count)

msprobe/test/pytorch_ut/hook_module/test_api_registry.py

@@ -0,0 +1,130 @@
+import unittest
+from msprobe.pytorch.hook_module.api_registry import ApiRegistry, torch_version_above_2, is_gpu
+
+
+class TestApiRegistry(unittest.TestCase):
+    
+    def test_store_ori_attr(self):
+        class A():
+            a1 = 1
+        class B():
+            a = A()
+            b1 = 1
+            b2 = 2
+        
+        api_list = ["a.a1", "b1", "b2"]
+        expect_output = {"a.a1":1, "b1":1, "b2":2}
+        actual_output = dict()
+        ApiRegistry.store_ori_attr(B, api_list, actual_output)
+        self.assertEqual(actual_output, expect_output)
+
+    
+    def test_set_api_attr(self):
+        class A():
+            a1 = 1
+        class B():
+            a = A().__class__
+            b1 = 1
+        
+        attr_dict = {"a.a2":2, "b2":2, "b3":3}
+        ApiRegistry.set_api_attr(B, attr_dict)
+
+        for k, v in attr_dict.items():
+            if '.' in k:
+                sub_module_name, sub_op = k.rsplit('.', 1)
+                sub_module = getattr(B, sub_module_name, None)
+                
+                self.assertEqual(getattr(sub_module, sub_op), v)
+            else:
+                self.assertEqual(getattr(B, k), v)
+    
+    def test_api_modularity(self):
+
+        import torch
+        import torch.distributed as dist
+        #import torch_npu   #门禁没有安装torch_npu
+        from msprobe.pytorch.hook_module.api_registry import torch_without_guard_version, npu_distributed_api, is_gpu, torch_version_above_2
+
+        
+
+        reg = ApiRegistry()
+        attr_dict = {"b2":2, "b3":3}
+        reg.tensor_hook_attr = attr_dict
+        reg.torch_hook_attr = attr_dict
+        reg.functional_hook_attr = attr_dict
+        reg.distributed_hook_attr = attr_dict
+        reg.npu_distributed_hook_attr = attr_dict
+        reg.aten_hook_attr = attr_dict
+        reg.vf_hook_attr = attr_dict
+        reg.torch_npu_hook_attr = attr_dict
+
+        reg.api_modularity()
+        self.assertEqual(torch.Tensor.b2, 2)
+
+        self.assertEqual(torch.b2, 2)
+        self.assertEqual(torch.nn.functional.b2, 2)
+        self.assertEqual(dist.b2, 2)
+        self.assertEqual(dist.distributed_c10d.b2, 2)
+        #if not is_gpu and not torch_without_guard_version:
+            #self.assertEqual(torch_npu.distributed.b2, 2)
+            #self.assertEqual(torch_npu.distributed.distributed_c10d.b2, 2)
+        if torch_version_above_2:
+            self.assertEqual(torch.ops.aten.b2, 2)
+        self.assertEqual(torch._VF.b2, 2)
+        #if not is_gpu:
+            #self.assertEqual(torch_npu.b2, 2)
+    
+
+    def test_api_originality(self):
+        import torch
+        import torch.distributed as dist
+        #import torch_npu      #门禁没有安装torch_npu
+        from msprobe.pytorch.hook_module.api_registry import torch_without_guard_version, npu_distributed_api, is_gpu, torch_version_above_2
+
+        
+
+        reg = ApiRegistry()
+        attr_dict = {"b2":2, "b3":3}
+        reg.tensor_hook_attr = attr_dict
+        reg.torch_hook_attr = attr_dict
+        reg.functional_hook_attr = attr_dict
+        reg.distributed_hook_attr = attr_dict
+        reg.npu_distributed_hook_attr = attr_dict
+        reg.aten_hook_attr = attr_dict
+        reg.vf_hook_attr = attr_dict
+        reg.torch_npu_hook_attr = attr_dict
+
+        reg.api_originality()
+        self.assertEqual(torch.Tensor.b2, 2)
+
+        self.assertEqual(torch.b2, 2)
+        self.assertEqual(torch.nn.functional.b2, 2)
+        self.assertEqual(dist.b2, 2)
+        self.assertEqual(dist.distributed_c10d.b2, 2)
+        #if not is_gpu and not torch_without_guard_version:
+            #self.assertEqual(torch_npu.distributed.b2, 2)
+            #self.assertEqual(torch_npu.distributed.distributed_c10d.b2, 2)
+        if torch_version_above_2:
+            self.assertEqual(torch.ops.aten.b2, 2)
+        self.assertEqual(torch._VF.b2, 2)
+        #if not is_gpu:
+            #self.assertEqual(torch_npu.b2, 2)
+
+    def test_initialize_hook(self):
+        def hook_test():
+            pass
+
+        reg = ApiRegistry()
+        reg.initialize_hook(hook_test)
+        empty_list = []
+        self.assertFalse(empty_list==reg.tensor_hook_attr)
+        self.assertFalse(empty_list==reg.torch_hook_attr)
+        self.assertFalse(empty_list==reg.functional_hook_attr)
+        self.assertFalse(empty_list==reg.distributed_hook_attr)
+        self.assertFalse(empty_list==reg.npu_distributed_hook_attr)
+        if torch_version_above_2:
+            #print(True)
+            self.assertFalse(empty_list==reg.aten_hook_attr)
+        if not is_gpu:
+            #print(True)
+            self.assertFalse(empty_list==reg.torch_npu_hook_attr)

msprobe/test/pytorch_ut/hook_module/test_hook_module.py

@@ -0,0 +1,42 @@
+import unittest
+from unittest.mock import patch, Mock
+
+from msprobe.pytorch.hook_module.hook_module import HOOKModule
+
+class TestHookModule(unittest.TestCase):
+    def test_call_1(self):
+        def forward_pre_hook():
+            return "result_input", "result_kwargs"
+        def forward_hook():
+            return 2
+        def backward_hook():
+            pass
+
+        def hook(prefix):
+            return forward_pre_hook, forward_hook, backward_hook
+        HOOKModule.prefix_op_name_ = "123"
+        test = HOOKModule(hook)
+        test._call_func = Mock(return_value=1)
+        result = test()
+        self.assertEqual(result, 1)
+
+    def test_call_2(self):
+        def forward_pre_hook(nope, input, kwargs):
+            return input, kwargs
+        def forward_hook(nope, input, kwargs, result):
+            return input
+        def backward_hook():
+            pass
+
+        def hook(prefix):
+            return forward_pre_hook, forward_hook, backward_hook
+        HOOKModule.prefix_op_name_ = "123"
+        input = 2
+        test = HOOKModule(hook)
+
+        def temp_forward(*input, **kwargs):
+            return input
+
+        test.forward = Mock(return_value=1)
+        result = test(input)
+        self.assertEqual(result, (input, ))

msprobe/test/pytorch_ut/hook_module/test_wrap_aten.py

@@ -0,0 +1,65 @@
+import unittest
+import torch
+from msprobe.pytorch.hook_module.wrap_aten import AtenOPTemplate, AtenOPPacketTemplate
+
+
+def hook(name):
+    def forward_pre_hook(nope, input, kwargs):
+            return input, kwargs
+    def forward_hook(nope, input, kwargs, result):
+            return 2
+    def backward_hook():
+            pass
+    
+    return forward_pre_hook, forward_hook, backward_hook
+
+
+
+class TestWrapAten(unittest.TestCase):
+    def setUp(self):
+        self.aten_op = AtenOPPacketTemplate(torch.ops.aten.convolution, hook)
+        
+    def test_atenop_attribute(self):
+        if torch.__version__.split("+")[0] <= '2.0':
+            return
+        self.setUp()
+        self.assertEqual(self.aten_op.default.op, torch.ops.aten.convolution.default)
+        self.assertEqual(self.aten_op.out.op, torch.ops.aten.convolution.out)
+
+    def test_atenop_forward(self):
+        if torch.__version__.split("+")[0] <= '2.0':
+            return
+        self.setUp()
+        image = torch.randn(4, 3, 24, 24)
+        kernel = torch.randn(10, 3, 3, 3)
+        functional_out = torch.nn.functional.conv2d(image, kernel, stride=[1, 1],
+                                                    padding=[1, 1], dilation=[1, 1], groups=1, bias=None)
+        aten_out = self.aten_op(image, kernel, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1)
+        self.assertTrue(aten_out == 2)
+
+    def test_atenop_overload_forward(self):
+        if torch.__version__.split("+")[0] <= '2.0':
+            return
+        self.setUp()
+        image = torch.randn(4, 3, 24, 24)
+        kernel = torch.randn(10, 3, 3, 3)
+        functional_out = torch.nn.functional.conv2d(image, kernel, stride=[1, 1],
+                                                    padding=[1, 1], dilation=[1, 1], groups=1, bias=None)
+        aten_out = self.aten_op(image, kernel, None, [1, 1], [1, 1], [1, 1], False, [0, 0], 1)
+        self.assertTrue(aten_out == 2)
+
+    def test_atenop_nonattr(self):
+        if torch.__version__.split("+")[0] <= '2.0':
+            return
+        self.setUp()
+        self.assertRaises(AttributeError, getattr, self.aten_op, "foo")
+
+    def test_atenop_overloads(self):
+        if torch.__version__.split("+")[0] <= '2.0':
+            return
+        self.setUp()
+        self.assertEqual(self.aten_op.overloads(), self.aten_op.opPacket.overloads())
+
+
+
+