已合并
test(data): adapt torch.utils.data tests on NPU #37377
yulin520创建于 6月1日
test(data): adapt torch.utils.data tests on NPU #37377
已合并
共 2 个文件变更+293-5
| @@ -1,16 +1,171 @@ | |||
| 1 | diff --git a/test/test_dataloader.py b/test/test_dataloader.py | 1 | diff --git a/test/test_dataloader.py b/test/test_dataloader.py |
| 2 | -index b2f47c4..f69a675 100644 | 2 | +index b2f47c4..4e687a3 100644 |
| 3 | --- a/test/test_dataloader.py | 3 | --- a/test/test_dataloader.py |
| 4 | +++ b/test/test_dataloader.py | 4 | +++ b/test/test_dataloader.py |
| 5 | -@@ -18,6 +18,7 @@ import unittest | 5 | +@@ -18,6 +18,8 @@ import unittest |
| 6 | import warnings | 6 | import warnings |
| 7 | 7 | ||
| 8 | import torch | 8 | import torch |
| 9 | +import torch_npu | 9 | +import torch_npu |
| 10 | ++from torch_npu.contrib import transfer_to_npu | ||
| 10 | import torch.utils.data.datapipes as dp | 11 | import torch.utils.data.datapipes as dp |
| 11 | from torch import multiprocessing as mp | 12 | from torch import multiprocessing as mp |
| 12 | from torch._utils import ExceptionWrapper | 13 | from torch._utils import ExceptionWrapper |
| 13 | -@@ -1197,6 +1198,11 @@ class TestDataLoader(TestCase): | 14 | +@@ -392,37 +394,37 @@ class CountingIterableDataset(IterableDataset): |
| 15 | + ) | ||
| 16 | + class TestTensorDataset(TestCase): | ||
| 17 | + def test_len(self): | ||
| 18 | +- source = TensorDataset(torch.randn(15, 10, 2, 3, 4, 5), torch.randperm(15)) | ||
| 19 | ++ source = TensorDataset(torch.randn(15, 10, 2, 3, 4, 5).npu(), torch.randperm(15).npu()) | ||
| 20 | + self.assertEqual(len(source), 15) | ||
| 21 | + | ||
| 22 | + def test_getitem(self): | ||
| 23 | +- t = torch.randn(15, 10, 2, 3, 4, 5) | ||
| 24 | +- l = torch.randn(15, 10) | ||
| 25 | ++ t = torch.randn(15, 10, 2, 3, 4, 5).npu() | ||
| 26 | ++ l = torch.randn(15, 10).npu() | ||
| 27 | + source = TensorDataset(t, l) | ||
| 28 | + for i in range(15): | ||
| 29 | + self.assertEqual(t[i], source[i][0]) | ||
| 30 | + self.assertEqual(l[i], source[i][1]) | ||
| 31 | + | ||
| 32 | + def test_getitem_1d(self): | ||
| 33 | +- t = torch.randn(15) | ||
| 34 | +- l = torch.randn(15) | ||
| 35 | ++ t = torch.randn(15).npu() | ||
| 36 | ++ l = torch.randn(15).npu() | ||
| 37 | + source = TensorDataset(t, l) | ||
| 38 | + for i in range(15): | ||
| 39 | + self.assertEqual(t[i], source[i][0]) | ||
| 40 | + self.assertEqual(l[i], source[i][1]) | ||
| 41 | + | ||
| 42 | + def test_single_tensor(self): | ||
| 43 | +- t = torch.randn(5, 10) | ||
| 44 | ++ t = torch.randn(5, 10).npu() | ||
| 45 | + source = TensorDataset(t) | ||
| 46 | + self.assertEqual(len(source), 5) | ||
| 47 | + for i in range(5): | ||
| 48 | + self.assertEqual(t[i], source[i][0]) | ||
| 49 | + | ||
| 50 | + def test_many_tensors(self): | ||
| 51 | +- t0 = torch.randn(5, 10, 2, 3, 4, 5) | ||
| 52 | +- t1 = torch.randn(5, 10) | ||
| 53 | +- t2 = torch.randn(5, 10, 2, 5) | ||
| 54 | +- t3 = torch.randn(5, 10, 3, 7) | ||
| 55 | ++ t0 = torch.randn(5, 10, 2, 3, 4, 5).npu() | ||
| 56 | ++ t1 = torch.randn(5, 10).npu() | ||
| 57 | ++ t2 = torch.randn(5, 10, 2, 5).npu() | ||
| 58 | ++ t3 = torch.randn(5, 10, 3, 7).npu() | ||
| 59 | + source = TensorDataset(t0, t1, t2, t3) | ||
| 60 | + self.assertEqual(len(source), 5) | ||
| 61 | + for i in range(5): | ||
| 62 | +@@ -447,36 +449,36 @@ class TestStackDataset(TestCase): | ||
| 63 | + def test_mixed(self): | ||
| 64 | + with self.assertRaisesRegex(ValueError, "Supported either"): | ||
| 65 | + StackDataset( | ||
| 66 | +- TensorDataset(torch.randn(15, 10)), a=TensorDataset(torch.randn(10, 15)) | ||
| 67 | ++ TensorDataset(torch.randn(15, 10).npu()), a=TensorDataset(torch.randn(10, 15).npu()) | ||
| 68 | + ) | ||
| 69 | + | ||
| 70 | + def test_size_mismatch(self): | ||
| 71 | + with self.assertRaisesRegex(ValueError, "Size mismatch between datasets"): | ||
| 72 | + StackDataset( | ||
| 73 | +- TensorDataset(torch.randn(15, 10)), TensorDataset(torch.randn(10, 15)) | ||
| 74 | ++ TensorDataset(torch.randn(15, 10).npu()), TensorDataset(torch.randn(10, 15).npu()) | ||
| 75 | + ) | ||
| 76 | + with self.assertRaisesRegex(ValueError, "Size mismatch between datasets"): | ||
| 77 | + StackDataset( | ||
| 78 | +- a=TensorDataset(torch.randn(15, 10)), | ||
| 79 | +- b=TensorDataset(torch.randn(10, 15)), | ||
| 80 | ++ a=TensorDataset(torch.randn(15, 10).npu()), | ||
| 81 | ++ b=TensorDataset(torch.randn(10, 15).npu()), | ||
| 82 | + ) | ||
| 83 | + | ||
| 84 | + def test_len(self): | ||
| 85 | + source = StackDataset( | ||
| 86 | +- TensorDataset(torch.randn(15, 10)), TensorDataset(torch.randn(15)) | ||
| 87 | ++ TensorDataset(torch.randn(15, 10).npu()), TensorDataset(torch.randn(15).npu()) | ||
| 88 | + ) | ||
| 89 | + self.assertEqual(len(source), 15) | ||
| 90 | +- source = StackDataset(TensorDataset(torch.randn(15, 10))) | ||
| 91 | ++ source = StackDataset(TensorDataset(torch.randn(15, 10).npu())) | ||
| 92 | + self.assertEqual(len(source), 15) | ||
| 93 | + source = StackDataset( | ||
| 94 | +- a=TensorDataset(torch.randn(15, 10)), b=TensorDataset(torch.randn(15)) | ||
| 95 | ++ a=TensorDataset(torch.randn(15, 10).npu()), b=TensorDataset(torch.randn(15).npu()) | ||
| 96 | + ) | ||
| 97 | + self.assertEqual(len(source), 15) | ||
| 98 | +- source = StackDataset(a=TensorDataset(torch.randn(15, 10))) | ||
| 99 | ++ source = StackDataset(a=TensorDataset(torch.randn(15, 10).npu())) | ||
| 100 | + self.assertEqual(len(source), 15) | ||
| 101 | + | ||
| 102 | + def test_single(self): | ||
| 103 | +- t = TensorDataset(torch.randn(15, 10)) | ||
| 104 | ++ t = TensorDataset(torch.randn(15, 10).npu()) | ||
| 105 | + source = StackDataset(t) | ||
| 106 | + for i in range(15): | ||
| 107 | + self.assertEqual(t[i], source[i][0]) | ||
| 108 | +@@ -485,8 +487,8 @@ class TestStackDataset(TestCase): | ||
| 109 | + self.assertEqual(t[i], source[i]["a"]) | ||
| 110 | + | ||
| 111 | + def test_getitem(self): | ||
| 112 | +- t = TensorDataset(torch.randn(15, 10)) | ||
| 113 | +- l = TensorDataset(torch.randn(15, 5, 4)) | ||
| 114 | ++ t = TensorDataset(torch.randn(15, 10).npu()) | ||
| 115 | ++ l = TensorDataset(torch.randn(15, 5, 4).npu()) | ||
| 116 | + source = StackDataset(t, l) | ||
| 117 | + for i in range(15): | ||
| 118 | + self.assertEqual(t[i], source[i][0]) | ||
| 119 | +@@ -499,7 +501,7 @@ class TestStackDataset(TestCase): | ||
| 120 | + def test_getitems(self): | ||
| 121 | + class GetItemsDataset(Dataset): | ||
| 122 | + def __init__(self) -> None: | ||
| 123 | +- self.data = torch.randn(4) | ||
| 124 | ++ self.data = torch.randn(4).npu() | ||
| 125 | + | ||
| 126 | + def __getitem__(self, item): | ||
| 127 | + return self.data[item] | ||
| 128 | +@@ -528,7 +530,7 @@ class TestStackDataset(TestCase): | ||
| 129 | + def test_getitems_raises_index_error(self): | ||
| 130 | + class GetItemsDataset(Dataset): | ||
| 131 | + def __init__(self) -> None: | ||
| 132 | +- self.data = torch.randn(4) | ||
| 133 | ++ self.data = torch.randn(4).npu() | ||
| 134 | + | ||
| 135 | + def __getitem__(self, item): | ||
| 136 | + return self.data[item] | ||
| 137 | +@@ -550,7 +552,7 @@ class TestStackDataset(TestCase): | ||
| 138 | + def test_getitems_value_error(self): | ||
| 139 | + class GetItemsDataset(Dataset): | ||
| 140 | + def __init__(self) -> None: | ||
| 141 | +- self.data = torch.randn(4) | ||
| 142 | ++ self.data = torch.randn(4).npu() | ||
| 143 | + | ||
| 144 | + def __getitem__(self, item): | ||
| 145 | + return self.data[item] | ||
| 146 | +@@ -604,9 +606,9 @@ class TestConcatDataset(TestCase): | ||
| 147 | + result[11] | ||
| 148 | + | ||
| 149 | + def test_add_dataset(self): | ||
| 150 | +- d1 = TensorDataset(torch.rand(7, 3, 28, 28), torch.rand(7)) | ||
| 151 | +- d2 = TensorDataset(torch.rand(7, 3, 28, 28), torch.rand(7)) | ||
| 152 | +- d3 = TensorDataset(torch.rand(7, 3, 28, 28), torch.rand(7)) | ||
| 153 | ++ d1 = TensorDataset(torch.rand(7, 3, 28, 28).npu(), torch.rand(7).npu()) | ||
| 154 | ++ d2 = TensorDataset(torch.rand(7, 3, 28, 28).npu(), torch.rand(7).npu()) | ||
| 155 | ++ d3 = TensorDataset(torch.rand(7, 3, 28, 28).npu(), torch.rand(7).npu()) | ||
| 156 | + result = d1 + d2 + d3 | ||
| 157 | + self.assertEqual(21, len(result)) | ||
| 158 | + self.assertEqual(0, (d1[0][0] - result[0][0]).abs().sum()) | ||
| 159 | +@@ -614,7 +616,7 @@ class TestConcatDataset(TestCase): | ||
| 160 | + self.assertEqual(0, (d3[0][0] - result[14][0]).abs().sum()) | ||
| 161 | + | ||
| 162 | + def test_iterable_dataset_err(self): | ||
| 163 | +- d1 = TensorDataset(torch.rand(7, 3, 28, 28), torch.rand(7)) | ||
| 164 | ++ d1 = TensorDataset(torch.rand(7, 3, 28, 28).npu(), torch.rand(7).npu()) | ||
| 165 | + it1 = CountingIterableDataset(5) | ||
| 166 | + it2 = CountingIterableDataset(10) | ||
| 167 | + | ||
| 168 | +@@ -1197,6 +1199,11 @@ class TestDataLoader(TestCase): | ||
| 14 | self.dataset = TensorDataset(self.data, self.labels) | 169 | self.dataset = TensorDataset(self.data, self.labels) |
| 15 | self.persistent_workers = False | 170 | self.persistent_workers = False |
| 16 | 171 | ||
| @@ -22,7 +177,7 @@ index b2f47c4..f69a675 100644 | |||
| 22 | def _get_data_loader(self, dataset, **kwargs): | 177 | def _get_data_loader(self, dataset, **kwargs): |
| 23 | persistent_workers = kwargs.get("persistent_workers", self.persistent_workers) | 178 | persistent_workers = kwargs.get("persistent_workers", self.persistent_workers) |
| 24 | if persistent_workers and kwargs.get("num_workers", 0) == 0: | 179 | if persistent_workers and kwargs.get("num_workers", 0) == 0: |
| 25 | -@@ -1340,9 +1346,11 @@ except RuntimeError as e: | 180 | +@@ -1340,9 +1347,11 @@ except RuntimeError as e: |
| 26 | self.assertRaises(ValueError, fn) | 181 | self.assertRaises(ValueError, fn) |
| 27 | 182 | ||
| 28 | def test_sequential_nonbatch(self): | 183 | def test_sequential_nonbatch(self): |
| @@ -34,7 +189,44 @@ index b2f47c4..f69a675 100644 | |||
| 34 | self._test_sequential(self._get_data_loader(self.dataset)) | 189 | self._test_sequential(self._get_data_loader(self.dataset)) |
| 35 | self._test_sequential(self._get_data_loader(self.dataset, batch_size=2)) | 190 | self._test_sequential(self._get_data_loader(self.dataset, batch_size=2)) |
| 36 | 191 | ||
| 37 | -@@ -2804,6 +2812,16 @@ except RuntimeError as e: | 192 | +@@ -1866,6 +1875,16 @@ except RuntimeError as e: |
| 193 | + list(iter(ChainDataset([dataset1, self.dataset]))) | ||
| 194 | + | ||
| 195 | + @unittest.skipIf(not TEST_CUDA_IPC, "CUDA IPC not available") | ||
| 196 | ++ @unittest.skipIf( | ||
| 197 | ++ torch.npu.is_available() and not any( | ||
| 198 | ++ conf.strip() == "expandable_segments:True" | ||
| 199 | ++ for conf in os.environ.get( | ||
| 200 | ++ "PYTORCH_NPU_ALLOC_CONF", "" | ||
| 201 | ++ ).split(",") | ||
| 202 | ++ ), | ||
| 203 | ++ "NPU IPC tensor sharing from DataLoader workers " | ||
| 204 | ++ "requires expandable_segments:True", | ||
| 205 | ++ ) | ||
| 206 | + @skipIfRocm # https://github.com/pytorch/pytorch/issues/90940 | ||
| 207 | + def test_multiprocessing_contexts(self): | ||
| 208 | + reference = [ | ||
| 209 | +@@ -2232,7 +2251,7 @@ except RuntimeError as e: | ||
| 210 | + def test_distributed_sampler_invalid_rank(self): | ||
| 211 | + from torch.utils.data.distributed import DistributedSampler | ||
| 212 | + | ||
| 213 | +- dataset = torch.IntTensor(range(10)) | ||
| 214 | ++ dataset = torch.IntTensor(range(10)).npu() | ||
| 215 | + with self.assertRaisesRegex(ValueError, "Invalid rank"): | ||
| 216 | + sampler = DistributedSampler(dataset, 3, 3) | ||
| 217 | + | ||
| 218 | +@@ -2244,8 +2263,8 @@ except RuntimeError as e: | ||
| 219 | + | ||
| 220 | + num_processes = 4 | ||
| 221 | + num_batches = 9 | ||
| 222 | +- data_set = torch.IntTensor(range(num_batches)) | ||
| 223 | +- scanned_data = torch.IntTensor([]) | ||
| 224 | ++ data_set = torch.IntTensor(range(num_batches)).npu() | ||
| 225 | ++ scanned_data = torch.IntTensor([]).npu() | ||
| 226 | + for i in range(num_processes): | ||
| 227 | + s = DistributedSampler(data_set, num_processes, i) | ||
| 228 | + d_loader = self._get_data_loader( | ||
| 229 | +@@ -2804,6 +2823,16 @@ except RuntimeError as e: | ||
| 38 | self.assertEqual(collated, torch.tensor(arr)) | 230 | self.assertEqual(collated, torch.tensor(arr)) |
| 39 | self.assertEqual(collated.dtype, torch.bool) | 231 | self.assertEqual(collated.dtype, torch.bool) |
| 40 | 232 | ||
| @@ -0,0 +1,96 @@ | |||
| 1 | +diff --git a/test/test_datapipe.py b/test/test_datapipe.py | ||
| 2 | +index 2a57bef..bb8bbb0 100644 | ||
| 3 | +--- a/test/test_datapipe.py | ||
| 4 | ++++ b/test/test_datapipe.py | ||
| 5 | +@@ -27,6 +27,8 @@ from unittest import skipIf | ||
| 6 | + import numpy as np | ||
| 7 | + | ||
| 8 | + import torch | ||
| 9 | ++import torch_npu | ||
| 10 | ++from torch_npu.contrib import transfer_to_npu | ||
| 11 | + import torch.nn as nn | ||
| 12 | + import torch.utils.data.datapipes as dp | ||
| 13 | + import torch.utils.data.graph | ||
| 14 | +@@ -1387,19 +1389,19 @@ class TestFunctionalIterDataPipe(TestCase): | ||
| 15 | + input_dp = dp.iter.IterableWrapper(range(target_length)) | ||
| 16 | + | ||
| 17 | + def fn(item, dtype=torch.float, *, sum=False): | ||
| 18 | +- data = torch.tensor(item, dtype=dtype) | ||
| 19 | ++ data = torch.tensor(item, dtype=dtype).npu() | ||
| 20 | + return data if not sum else data.sum() | ||
| 21 | + | ||
| 22 | + # Functional Test: apply to each element correctly | ||
| 23 | + map_dp = input_dp.map(fn) | ||
| 24 | + self.assertEqual(target_length, len(map_dp)) | ||
| 25 | + for x, y in zip(map_dp, range(target_length)): | ||
| 26 | +- self.assertEqual(x, torch.tensor(y, dtype=torch.float)) | ||
| 27 | ++ self.assertEqual(x, torch.tensor(y, dtype=torch.float).npu()) | ||
| 28 | + | ||
| 29 | + # Functional Test: works with partial function | ||
| 30 | + map_dp = input_dp.map(partial(fn, dtype=torch.int, sum=True)) | ||
| 31 | + for x, y in zip(map_dp, range(target_length)): | ||
| 32 | +- self.assertEqual(x, torch.tensor(y, dtype=torch.int).sum()) | ||
| 33 | ++ self.assertEqual(x, torch.tensor(y, dtype=torch.int).npu().sum()) | ||
| 34 | + | ||
| 35 | + # __len__ Test: inherits length from source DataPipe | ||
| 36 | + self.assertEqual(target_length, len(map_dp)) | ||
| 37 | +@@ -1728,7 +1730,7 @@ class TestFunctionalIterDataPipe(TestCase): | ||
| 38 | + input_dp = dp.iter.IterableWrapper(arrs) | ||
| 39 | + | ||
| 40 | + def _collate_fn(batch, default_type=torch.float): | ||
| 41 | +- return torch.tensor(sum(batch), dtype=default_type) | ||
| 42 | ++ return torch.tensor(sum(batch), dtype=default_type).npu() | ||
| 43 | + | ||
| 44 | + # Functional Test: defaults to the default collate function when a custom one is not specified | ||
| 45 | + collate_dp = input_dp.collate() | ||
| 46 | +@@ -1738,21 +1740,21 @@ class TestFunctionalIterDataPipe(TestCase): | ||
| 47 | + # Functional Test: custom collate function | ||
| 48 | + collate_dp = input_dp.collate(collate_fn=_collate_fn) | ||
| 49 | + for x, y in zip(arrs, collate_dp): | ||
| 50 | +- self.assertEqual(torch.tensor(sum(x), dtype=torch.float), y) | ||
| 51 | ++ self.assertEqual(torch.tensor(sum(x), dtype=torch.float).npu(), y) | ||
| 52 | + | ||
| 53 | + # Functional Test: custom, partial collate function | ||
| 54 | + collate_dp = input_dp.collate(partial(_collate_fn, default_type=torch.int)) | ||
| 55 | + for x, y in zip(arrs, collate_dp): | ||
| 56 | +- self.assertEqual(torch.tensor(sum(x), dtype=torch.int), y) | ||
| 57 | ++ self.assertEqual(torch.tensor(sum(x), dtype=torch.int).npu(), y) | ||
| 58 | + | ||
| 59 | + # Reset Test: reset the DataPipe and results are still correct | ||
| 60 | + n_elements_before_reset = 1 | ||
| 61 | + res_before_reset, res_after_reset = reset_after_n_next_calls( | ||
| 62 | + collate_dp, n_elements_before_reset | ||
| 63 | + ) | ||
| 64 | +- self.assertEqual([torch.tensor(6, dtype=torch.int)], res_before_reset) | ||
| 65 | ++ self.assertEqual([torch.tensor(6, dtype=torch.int).npu()], res_before_reset) | ||
| 66 | + for x, y in zip(arrs, res_after_reset): | ||
| 67 | +- self.assertEqual(torch.tensor(sum(x), dtype=torch.int), y) | ||
| 68 | ++ self.assertEqual(torch.tensor(sum(x), dtype=torch.int).npu(), y) | ||
| 69 | + | ||
| 70 | + # __len__ Test: __len__ is inherited | ||
| 71 | + self.assertEqual(len(input_dp), len(collate_dp)) | ||
| 72 | +@@ -2334,21 +2336,21 @@ class TestFunctionalMapDataPipe(TestCase): | ||
| 73 | + input_dp = dp.map.SequenceWrapper(arr) | ||
| 74 | + | ||
| 75 | + def fn(item, dtype=torch.float, *, sum=False): | ||
| 76 | +- data = torch.tensor(item, dtype=dtype) | ||
| 77 | ++ data = torch.tensor(item, dtype=dtype).npu() | ||
| 78 | + return data if not sum else data.sum() | ||
| 79 | + | ||
| 80 | + map_dp = input_dp.map(fn) | ||
| 81 | + self.assertEqual(len(input_dp), len(map_dp)) | ||
| 82 | + for index in arr: | ||
| 83 | + self.assertEqual( | ||
| 84 | +- map_dp[index], torch.tensor(input_dp[index], dtype=torch.float) | ||
| 85 | ++ map_dp[index], torch.tensor(input_dp[index], dtype=torch.float).npu() | ||
| 86 | + ) | ||
| 87 | + | ||
| 88 | + map_dp = input_dp.map(partial(fn, dtype=torch.int, sum=True)) | ||
| 89 | + self.assertEqual(len(input_dp), len(map_dp)) | ||
| 90 | + for index in arr: | ||
| 91 | + self.assertEqual( | ||
| 92 | +- map_dp[index], torch.tensor(input_dp[index], dtype=torch.int).sum() | ||
| 93 | ++ map_dp[index], torch.tensor(input_dp[index], dtype=torch.int).npu().sum() | ||
| 94 | + ) | ||
| 95 | + | ||
| 96 | + def test_batch_mapdatapipe(self): | ||