已合并
pta supports to create nested tensor #30196
culechan创建于 1月29日
pta supports to create nested tensor #30196
已合并
共 4 个文件变更+117-33
| @@ -109,18 +109,6 @@ class TestPtaUnsupportApi(TestCase): | |||
| 109 | traced_cell = torch.jit.trace(model, (x, h)) | 109 | traced_cell = torch.jit.trace(model, (x, h)) |
| 110 | traced_cell.bfloat16() | 110 | traced_cell.bfloat16() |
| 111 | 111 | ||
| 112 | - def test_nested_tensor_runtimeerror(self): | ||
| 113 | - with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."): | ||
| 114 | - a = torch.arange(3, dtype=torch.float).npu() | ||
| 115 | - b = torch.arange(3, dtype=torch.float).npu() | ||
| 116 | - torch.nested.nested_tensor([a, b]) | ||
| 117 | - | ||
| 118 | - def test_as_nested_tensor_runtimeerror(self): | ||
| 119 | - with self.assertRaisesRegex(RuntimeError, r"(.*) is not supported in npu."): | ||
| 120 | - a = torch.arange(3, dtype=torch.float).npu() | ||
| 121 | - b = torch.arange(3, dtype=torch.float).npu() | ||
| 122 | - torch.nested.as_nested_tensor([a, b]) | ||
| 123 | - | ||
| 124 | def test_Tensor_is_shared(self): | 112 | def test_Tensor_is_shared(self): |
| 125 | input_tensor = torch.tensor([1, 2, 3]) | 113 | input_tensor = torch.tensor([1, 2, 3]) |
| 126 | input_tensor.is_shared() | 114 | input_tensor.is_shared() |
| @@ -133,16 +121,6 @@ class TestPtaUnsupportApi(TestCase): | |||
| 133 | model = SimpleModel() | 121 | model = SimpleModel() |
| 134 | model.share_memory() | 122 | model.share_memory() |
| 135 | 123 | ||
| 136 | - def test_nested_tensor(self): | ||
| 137 | - a = torch.arange(3, dtype=torch.float) | ||
| 138 | - b = torch.arange(3, dtype=torch.float) | ||
| 139 | - torch.nested.nested_tensor([a, b]) | ||
| 140 | - | ||
| 141 | - def test_as_nested_tensor(self): | ||
| 142 | - a = torch.arange(3, dtype=torch.float) | ||
| 143 | - b = torch.arange(3, dtype=torch.float) | ||
| 144 | - torch.nested.as_nested_tensor([a, b]) | ||
| 145 | - | ||
| 146 | 124 | ||
| 147 | if __name__ == "__main__": | 125 | if __name__ == "__main__": |
| 148 | run_tests() | 126 | run_tests() |
| @@ -0,0 +1,116 @@ | |||
| 1 | +# Owner(s): ["module: nested tensor"] | ||
| 2 | + | ||
| 3 | +from collections import namedtuple, OrderedDict | ||
| 4 | +from multiprocessing.reduction import ForkingPickler | ||
| 5 | +import torch | ||
| 6 | +import numpy as np | ||
| 7 | +from torch.testing._internal.common_utils import parametrize, instantiate_parametrized_tests | ||
| 8 | +import torch_npu | ||
| 9 | +from torch_npu.testing.testcase import TestCase, run_tests | ||
| 10 | + | ||
| 11 | + | ||
| 12 | +class TestNestedTensor(TestCase): | ||
| 13 | + | ||
| 14 | + | ||
| 15 | + | ||
| 16 | + def test_2d_nested_tensor(self, batch_size, max_seq_len, vocab_size): | ||
| 17 | + data = [] | ||
| 18 | + nested_tensor_ref_list = [] | ||
| 19 | + for _ in range(batch_size): | ||
| 20 | + if max_seq_len == 0: | ||
| 21 | + length = 0 | ||
| 22 | + else: | ||
| 23 | + length = np.random.randint(1, max_seq_len) | ||
| 24 | + row = list(np.random.randint(low=0, high=vocab_size, size=(length,))) | ||
| 25 | + data.append(row) | ||
| 26 | + nested_tensor_ref_list.append(torch.tensor(row)) | ||
| 27 | + nested_tensor = torch.nested.nested_tensor(data) | ||
| 28 | + nested_tensor_list = nested_tensor.unbind() | ||
| 29 | + for i in range(batch_size): | ||
| 30 | + self.assertEqual(nested_tensor_list[i], nested_tensor_ref_list[i].type(torch.int64)) | ||
| 31 | + | ||
| 32 | + | ||
| 33 | + | ||
| 34 | + | ||
| 35 | + def test_3d_nested_tensor(self, batch_size, max_seq_len, vocab_size): | ||
| 36 | + data = [] | ||
| 37 | + nested_tensor_ref_list = [] | ||
| 38 | + for _ in range(batch_size): | ||
| 39 | + if max_seq_len == 0: | ||
| 40 | + length = 0 | ||
| 41 | + else: | ||
| 42 | + length = np.random.randint(1, max_seq_len) | ||
| 43 | + row = list(np.random.randint(low=0, high=vocab_size, size=(length,))) | ||
| 44 | + row = [list(item * np.arange(max_seq_len)) for item in row] | ||
| 45 | + data.append(row) | ||
| 46 | + nested_tensor_ref_list.append(torch.tensor(row)) | ||
| 47 | + nested_tensor = torch.nested.nested_tensor(data) | ||
| 48 | + nested_tensor_list = nested_tensor.unbind() | ||
| 49 | + for i in range(batch_size): | ||
| 50 | + self.assertEqual(nested_tensor_list[i], nested_tensor_ref_list[i].type(torch.int64)) | ||
| 51 | + | ||
| 52 | + | ||
| 53 | + | ||
| 54 | + | ||
| 55 | + def test_3d_nested_tensor_float(self, batch_size, max_seq_len, vocab_size): | ||
| 56 | + data = [] | ||
| 57 | + nested_tensor_ref_list = [] | ||
| 58 | + for _ in range(batch_size): | ||
| 59 | + if max_seq_len == 0: | ||
| 60 | + length = 0 | ||
| 61 | + else: | ||
| 62 | + length = np.random.randint(1, max_seq_len) | ||
| 63 | + row = list(np.random.randint(low=0, high=vocab_size, size=(length,))) | ||
| 64 | + row = [list(item * np.arange(max_seq_len)) for item in row] | ||
| 65 | + data.append(row) | ||
| 66 | + nested_tensor_ref_list.append(torch.tensor(row)) | ||
| 67 | + nested_tensor = torch.nested.nested_tensor(data) | ||
| 68 | + nested_tensor_list = nested_tensor.unbind() | ||
| 69 | + for i in range(batch_size): | ||
| 70 | + self.assertEqual(nested_tensor_list[i], nested_tensor_ref_list[i].type(torch.float32)) | ||
| 71 | + | ||
| 72 | + def _test_unbind_case(self, a, b): | ||
| 73 | + nt = torch.nested.nested_tensor([a, b], dtype=a.dtype) | ||
| 74 | + nt_list = nt.unbind() | ||
| 75 | + self.assertEqual(len(nt_list), 2) | ||
| 76 | + self.assertEqual(nt_list[0], a) | ||
| 77 | + self.assertEqual(nt_list[1], b) | ||
| 78 | + | ||
| 79 | + def test_unbind_case1(self): | ||
| 80 | + a = torch.tensor([[1, 2, 3], [4, 5, 6]]) | ||
| 81 | + b = torch.tensor([[7, 8], [10, 11]]) | ||
| 82 | + self._test_unbind_case(a, b) | ||
| 83 | + | ||
| 84 | + def test_unbind_case2(self): | ||
| 85 | + a = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32) | ||
| 86 | + b = torch.tensor([[7, 8], [10, 11]], dtype=torch.float32) | ||
| 87 | + self._test_unbind_case(a, b) | ||
| 88 | + | ||
| 89 | + def test_unbind_case3(self): | ||
| 90 | + a = torch.tensor([[], []]) | ||
| 91 | + b = torch.tensor([[], [], []]) | ||
| 92 | + self._test_unbind_case(a, b) | ||
| 93 | + | ||
| 94 | + def test_default_options_nested_tensor(self): | ||
| 95 | + default_nested_tensor = torch.nested.nested_tensor([]) | ||
| 96 | + default_tensor = torch.tensor([]) | ||
| 97 | + self.assertEqual(default_nested_tensor.dtype, default_tensor.dtype) | ||
| 98 | + self.assertEqual(default_nested_tensor.device, default_tensor.device) | ||
| 99 | + self.assertEqual(default_nested_tensor.layout, default_tensor.layout) | ||
| 100 | + self.assertEqual(default_nested_tensor.dim(), default_tensor.dim()) | ||
| 101 | + self.assertEqual(default_nested_tensor.requires_grad, default_tensor.requires_grad) | ||
| 102 | + | ||
| 103 | + def test_nested_tensor_size(self): | ||
| 104 | + nt = torch.nested.nested_tensor([torch.tensor([[1, 2, 3], [4, 5, 6]]), torch.tensor([[7, 8], [10, 11], [12, 13]])]) | ||
| 105 | + self.assertEqual(nt.size(0), 2) | ||
| 106 | + self.assertRaisesRegex(RuntimeError, | ||
| 107 | + "Given dimension 1 is irregular and does not have a size", | ||
| 108 | + lambda: nt.size(1), | ||
| 109 | + ) | ||
| 110 | + | ||
| 111 | + nt = torch.nested.nested_tensor([2]) | ||
| 112 | + self.assertEqual(nt.size(0), 1) | ||
| 113 | + | ||
| 114 | +if __name__ == '__main__': | ||
| 115 | + instantiate_parametrized_tests(TestNestedTensor) | ||
| 116 | + run_tests() | ||
| @@ -6,7 +6,7 @@ from functools import wraps | |||
| 6 | import torch | 6 | import torch |
| 7 | import torch_npu | 7 | import torch_npu |
| 8 | from torch_npu.utils._error_code import ErrCode, pta_error | 8 | from torch_npu.utils._error_code import ErrCode, pta_error |
| 9 | -from .unsupport_api import unsupported_Tensor_api, unsupported_nn_api, unsupported_nested_api | 9 | +from .unsupport_api import unsupported_Tensor_api, unsupported_nn_api |
| 10 | from .collect_env import get_cann_version | 10 | from .collect_env import get_cann_version |
| 11 | 11 | ||
| 12 | 12 | ||
| @@ -93,10 +93,6 @@ def _is_module_parameters_supported(*args, **kwargs): | |||
| 93 | return any(p.device is not None and p.device.type == "npu" for p in module_parameters) | 93 | return any(p.device is not None and p.device.type == "npu" for p in module_parameters) |
| 94 | 94 | ||
| 95 | 95 | ||
| 96 | -def _is_nested_tensor_npu_supported(*args, **kwargs): | ||
| 97 | - return any(torch.is_tensor(t) and t.is_npu for t in args[0]) | ||
| 98 | - | ||
| 99 | - | ||
| 100 | def _apply_wrap_func_to_modules(wrap_func, unsupported_modules): | 96 | def _apply_wrap_func_to_modules(wrap_func, unsupported_modules): |
| 101 | for attr_name, parent_module in unsupported_modules.items(): | 97 | for attr_name, parent_module in unsupported_modules.items(): |
| 102 | setattr(parent_module, attr_name, wrap_func(getattr(parent_module, attr_name))) | 98 | setattr(parent_module, attr_name, wrap_func(getattr(parent_module, attr_name))) |
| @@ -106,4 +102,3 @@ def _apply_wrap_func_to_modules(wrap_func, unsupported_modules): | |||
| 106 | def _add_intercept_methods(): | 102 | def _add_intercept_methods(): |
| 107 | _apply_wrap_func_to_modules(_create_wrap_func(_is_tensor_npu_supported), unsupported_Tensor_api) | 103 | _apply_wrap_func_to_modules(_create_wrap_func(_is_tensor_npu_supported), unsupported_Tensor_api) |
| 108 | _apply_wrap_func_to_modules(_create_wrap_func(_is_module_parameters_supported), unsupported_nn_api) | 104 | _apply_wrap_func_to_modules(_create_wrap_func(_is_module_parameters_supported), unsupported_nn_api) |
| 109 | - _apply_wrap_func_to_modules(_create_wrap_func(_is_nested_tensor_npu_supported), unsupported_nested_api) | ||
| @@ -16,8 +16,3 @@ unsupported_nn_api = { | |||
| 16 | "register_parameter": torch.jit.ScriptModule, | 16 | "register_parameter": torch.jit.ScriptModule, |
| 17 | "register_module": torch.jit.ScriptModule | 17 | "register_module": torch.jit.ScriptModule |
| 18 | } | 18 | } |
| 19 | - | ||
| 20 | -unsupported_nested_api = { | ||
| 21 | - "nested_tensor": torch.nested, | ||
| 22 | - "as_nested_tensor": torch.nested | ||
| 23 | -} | ||