已合并
test(nn): adapt PackedSequence tests on NPU #37115
Jinfan Liu创建于 5月29日
test(nn): adapt PackedSequence tests on NPU #37115
已合并
共 2 个文件变更+204-0
| @@ -0,0 +1,83 @@ | |||
| 1 | +# Owner(s): ["module: nn"] | ||
| 2 | + | ||
| 3 | +""" | ||
| 4 | +Add validation cases for torch.nn.utils.rnn.PackedSequence APIs on NPU: | ||
| 5 | +1. PyTorch community lacks direct validations for PackedSequence.count, | ||
| 6 | + PackedSequence.index, and PackedSequence.is_pinned. | ||
| 7 | +2. This file validates torch.nn.utils.rnn.PackedSequence, | ||
| 8 | + torch.nn.utils.rnn.PackedSequence.count, | ||
| 9 | + torch.nn.utils.rnn.PackedSequence.index, and | ||
| 10 | + torch.nn.utils.rnn.PackedSequence.is_pinned (extendable). | ||
| 11 | +""" | ||
| 12 | + | ||
| 13 | +import torch | ||
| 14 | +import torch.nn.utils.rnn as rnn_utils | ||
| 15 | +from torch.testing._internal.common_utils import TestCase, run_tests | ||
| 16 | + | ||
| 17 | + | ||
| 18 | +device_type = "npu" if hasattr(torch, "npu") and torch.npu.is_available() else "cpu" | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +class TestPackedSequenceAPIs(TestCase): | ||
| 22 | + | ||
| 23 | + def test_packed_sequence_constructor_and_to_on_npu(self): | ||
| 24 | + data = torch.randn(5, 10, device=device_type) | ||
| 25 | + batch_sizes = torch.tensor([3, 2], dtype=torch.int64) | ||
| 26 | + sorted_indices = torch.tensor([2, 0, 1], dtype=torch.int64, device=device_type) | ||
| 27 | + unsorted_indices = torch.tensor( | ||
| 28 | + [1, 2, 0], dtype=torch.int64, device=device_type | ||
| 29 | + ) | ||
| 30 | + | ||
| 31 | + packed = rnn_utils.PackedSequence( | ||
| 32 | + data, batch_sizes, sorted_indices, unsorted_indices | ||
| 33 | + ) | ||
| 34 | + | ||
| 35 | + self.assertIsInstance(packed, rnn_utils.PackedSequence) | ||
| 36 | + self.assertEqual(packed.data.device.type, device_type) | ||
| 37 | + self.assertEqual(packed.data.shape, torch.Size([5, 10])) | ||
| 38 | + self.assertEqual(packed.batch_sizes.device.type, "cpu") | ||
| 39 | + self.assertEqual(packed.sorted_indices.device.type, device_type) | ||
| 40 | + self.assertEqual(packed.unsorted_indices.device.type, device_type) | ||
| 41 | + self.assertFalse(packed.is_pinned()) | ||
| 42 | + | ||
| 43 | + packed_cpu = packed.to("cpu") | ||
| 44 | + self.assertEqual(packed_cpu.data.device.type, "cpu") | ||
| 45 | + self.assertEqual(packed_cpu.batch_sizes.device.type, "cpu") | ||
| 46 | + self.assertEqual(packed_cpu.sorted_indices.device.type, "cpu") | ||
| 47 | + self.assertEqual(packed_cpu.unsorted_indices.device.type, "cpu") | ||
| 48 | + | ||
| 49 | + packed_accelerator = packed_cpu.to(device_type) | ||
| 50 | + self.assertEqual(packed_accelerator.data.device.type, device_type) | ||
| 51 | + self.assertEqual(packed_accelerator.batch_sizes.device.type, "cpu") | ||
| 52 | + self.assertEqual(packed_accelerator.sorted_indices.device.type, device_type) | ||
| 53 | + self.assertEqual(packed_accelerator.unsorted_indices.device.type, device_type) | ||
| 54 | + | ||
| 55 | + def test_packed_sequence_namedtuple_methods_on_optional_indices(self): | ||
| 56 | + data = torch.tensor([1.0, 2.0, 3.0], device=device_type) | ||
| 57 | + batch_sizes = torch.tensor([2, 1], dtype=torch.int64) | ||
| 58 | + packed = rnn_utils.PackedSequence(data, batch_sizes) | ||
| 59 | + | ||
| 60 | + self.assertIsNone(packed.sorted_indices) | ||
| 61 | + self.assertIsNone(packed.unsorted_indices) | ||
| 62 | + self.assertEqual(packed.count(None), 2) | ||
| 63 | + self.assertEqual(packed.index(None), 2) | ||
| 64 | + self.assertEqual(packed.count("non_existent"), 0) | ||
| 65 | + with self.assertRaises(ValueError): | ||
| 66 | + packed.index("non_existent") | ||
| 67 | + | ||
| 68 | + def test_packed_sequence_rejects_accelerator_batch_sizes(self): | ||
| 69 | + data = torch.tensor([1.0, 2.0], device=device_type) | ||
| 70 | + batch_sizes = torch.tensor([2], dtype=torch.int64, device=device_type) | ||
| 71 | + | ||
| 72 | + if device_type == "cpu": | ||
| 73 | + packed = rnn_utils.PackedSequence(data, batch_sizes) | ||
| 74 | + self.assertEqual(packed.batch_sizes.device.type, "cpu") | ||
| 75 | + else: | ||
| 76 | + with self.assertRaisesRegex( | ||
| 77 | + ValueError, "batch_sizes should always be on CPU" | ||
| 78 | + ): | ||
| 79 | + rnn_utils.PackedSequence(data, batch_sizes) | ||
| 80 | + | ||
| 81 | + | ||
| 82 | +if __name__ == "__main__": | ||
| 83 | + run_tests() | ||
| @@ -0,0 +1,121 @@ | |||
| 1 | +diff --git a/test/nn/test_packed_sequence.py b/test/nn/test_packed_sequence.py | ||
| 2 | +index 016a5ef..3d57794 100644 | ||
| 3 | +--- a/test/nn/test_packed_sequence.py | ||
| 4 | ++++ b/test/nn/test_packed_sequence.py | ||
| 5 | + import torch | ||
| 6 | + import torch.nn.utils.rnn as rnn_utils | ||
| 7 | + from torch.testing._internal.common_utils import run_tests, TestCase | ||
| 8 | + | ||
| 9 | ++import torch_npu | ||
| 10 | ++from torch_npu.contrib import transfer_to_npu | ||
| 11 | ++ | ||
| 12 | + | ||
| 13 | + class PackedSequenceTest(TestCase): | ||
| 14 | + _type_by_name = { | ||
| 15 | + class PackedSequenceTest(TestCase): | ||
| 16 | + | ||
| 17 | + def test_pad_sequence_with_tensor_sequences(self): | ||
| 18 | + seq_tuple_input = torch.nn.utils.rnn.pad_sequence( | ||
| 19 | +- (torch.tensor([[7, 6]]), torch.tensor([[-7, -1]])) | ||
| 20 | ++ (torch.tensor([[7, 6]]).npu(), torch.tensor([[-7, -1]]).npu()) | ||
| 21 | + ) | ||
| 22 | + seq_tensor_input = torch.nn.utils.rnn.pad_sequence( | ||
| 23 | +- torch.tensor([[[7, 6]], [[-7, -1]]]) | ||
| 24 | ++ torch.tensor([[[7, 6]], [[-7, -1]]]).npu() | ||
| 25 | + ) | ||
| 26 | + self.assertEqual(seq_tuple_input, seq_tensor_input) | ||
| 27 | + self.assertEqual(seq_tuple_input.shape, torch.Size([1, 2, 2])) | ||
| 28 | + class PackedSequenceTest(TestCase): | ||
| 29 | + ) | ||
| 30 | + | ||
| 31 | + # single dimensional | ||
| 32 | +- a = torch.tensor([1, 2, 3]) | ||
| 33 | +- b = torch.tensor([4, 5]) | ||
| 34 | +- c = torch.tensor([6]) | ||
| 35 | ++ a = torch.tensor([1, 2, 3]).npu() | ||
| 36 | ++ b = torch.tensor([4, 5]).npu() | ||
| 37 | ++ c = torch.tensor([6]).npu() | ||
| 38 | + | ||
| 39 | + # batch_first = true | ||
| 40 | +- expected = torch.tensor([[4, 5, 0], [1, 2, 3], [6, 0, 0]]) | ||
| 41 | ++ expected = torch.tensor([[4, 5, 0], [1, 2, 3], [6, 0, 0]]).npu() | ||
| 42 | + padded = rnn_utils.pad_sequence([b, a, c], True) | ||
| 43 | + self.assertEqual(padded, expected) | ||
| 44 | + | ||
| 45 | + class PackedSequenceTest(TestCase): | ||
| 46 | + self.assertEqual(padded, expected.transpose(0, 1)) | ||
| 47 | + | ||
| 48 | + # padding_side = "left", batch_first=True | ||
| 49 | +- expected = torch.tensor([[0, 4, 5], [1, 2, 3], [0, 0, 6]]) | ||
| 50 | ++ expected = torch.tensor([[0, 4, 5], [1, 2, 3], [0, 0, 6]]).npu() | ||
| 51 | + padded = rnn_utils.pad_sequence( | ||
| 52 | + [b, a, c], | ||
| 53 | + batch_first=True, | ||
| 54 | + class PackedSequenceTest(TestCase): | ||
| 55 | + self.assertEqual(padded, expected.transpose(0, 1)) | ||
| 56 | + | ||
| 57 | + # pad with non-zero value | ||
| 58 | +- expected = torch.tensor([[4, 5, 1], [1, 2, 3], [6, 1, 1]]) | ||
| 59 | ++ expected = torch.tensor([[4, 5, 1], [1, 2, 3], [6, 1, 1]]).npu() | ||
| 60 | + padded = rnn_utils.pad_sequence([b, a, c], True, 1) | ||
| 61 | + self.assertEqual(padded, expected) | ||
| 62 | + | ||
| 63 | + # Test pad sorted sequence | ||
| 64 | +- expected = torch.tensor([[1, 2, 3], [4, 5, 0], [6, 0, 0]]) | ||
| 65 | ++ expected = torch.tensor([[1, 2, 3], [4, 5, 0], [6, 0, 0]]).npu() | ||
| 66 | + padded = rnn_utils.pad_sequence([a, b, c], True) | ||
| 67 | + self.assertEqual(padded, expected) | ||
| 68 | + | ||
| 69 | + class PackedSequenceTest(TestCase): | ||
| 70 | + trailing_dims = [4] * num_dim | ||
| 71 | + for i in range(1, maxlen + 1): | ||
| 72 | + seq_len = i * i | ||
| 73 | +- sequences.append(torch.rand(seq_len, 5, *trailing_dims)) | ||
| 74 | ++ sequences.append(torch.rand(seq_len, 5, *trailing_dims).npu()) | ||
| 75 | + random.shuffle(sequences) | ||
| 76 | + # batch first = true | ||
| 77 | + expected = torch.stack([pad(seq, maxlen * maxlen) for seq in sequences]) | ||
| 78 | + class PackedSequenceTest(TestCase): | ||
| 79 | + | ||
| 80 | + def test_unpad_sequence(self): | ||
| 81 | + # single dimensional | ||
| 82 | +- a = torch.tensor([1, 2, 3]) | ||
| 83 | +- b = torch.tensor([4, 5]) | ||
| 84 | +- c = torch.tensor([6]) | ||
| 85 | ++ a = torch.tensor([1, 2, 3]).npu() | ||
| 86 | ++ b = torch.tensor([4, 5]).npu() | ||
| 87 | ++ c = torch.tensor([6]).npu() | ||
| 88 | + sequences = [a, b, c] | ||
| 89 | + | ||
| 90 | + lengths = torch.as_tensor([v.size(0) for v in sequences]) | ||
| 91 | + class PackedSequenceTest(TestCase): | ||
| 92 | + trailing_dims = [4] * num_dim | ||
| 93 | + for i in range(1, maxlen + 1): | ||
| 94 | + seq_len = i * i | ||
| 95 | +- sequences.append(torch.rand(seq_len, 5, *trailing_dims)) | ||
| 96 | ++ sequences.append(torch.rand(seq_len, 5, *trailing_dims).npu()) | ||
| 97 | + random.shuffle(sequences) | ||
| 98 | + | ||
| 99 | + lengths = torch.as_tensor([v.size(0) for v in sequences]) | ||
| 100 | + class PackedSequenceTest(TestCase): | ||
| 101 | + | ||
| 102 | + def test_unpack_sequence(self): | ||
| 103 | + # single dimensional | ||
| 104 | +- a = torch.tensor([1, 2, 3]) | ||
| 105 | +- b = torch.tensor([4, 5]) | ||
| 106 | +- c = torch.tensor([6]) | ||
| 107 | ++ a = torch.tensor([1, 2, 3]).npu() | ||
| 108 | ++ b = torch.tensor([4, 5]).npu() | ||
| 109 | ++ c = torch.tensor([6]).npu() | ||
| 110 | + sequences = [a, b, c] | ||
| 111 | + | ||
| 112 | + packed_sequences = rnn_utils.pack_sequence(sequences, enforce_sorted=False) | ||
| 113 | + class PackedSequenceTest(TestCase): | ||
| 114 | + trailing_dims = [4] * num_dim | ||
| 115 | + for i in range(1, maxlen + 1): | ||
| 116 | + seq_len = i * i | ||
| 117 | +- sequences.append(torch.rand(seq_len, 5, *trailing_dims)) | ||
| 118 | ++ sequences.append(torch.rand(seq_len, 5, *trailing_dims).npu()) | ||
| 119 | + random.shuffle(sequences) | ||
| 120 | + | ||
| 121 | + packed_sequences = rnn_utils.pack_sequence(sequences, enforce_sorted=False) | ||