| # Owner(s): ["module: nn"] | |
| import itertools | |
| import random | |
| import torch | |
| from torch.testing._internal.common_utils import TestCase, run_tests | |
| import torch.nn.utils.rnn as rnn_utils | |
| class PackedSequenceTest(TestCase): | |
| _type_by_name = { | |
| 'torch.DoubleTensor': (torch.DoubleTensor, 'double'), | |
| 'torch.FloatTensor': (torch.FloatTensor, 'float'), | |
| # We leave out `'torch.HalfTensor': (torch.HalfTensor, 'half'),` | |
| # because of an error in `pad_packed_sequence` | |
| # > AttributeError: 'torch.HalfTensor' object has no attribute 'fill_' | |
| 'torch.LongTensor': (torch.LongTensor, 'long'), | |
| 'torch.IntTensor': (torch.IntTensor, 'int'), | |
| 'torch.ShortTensor': (torch.ShortTensor, 'short'), | |
| 'torch.CharTensor': (torch.CharTensor, 'char'), | |
| 'torch.ByteTensor': (torch.ByteTensor, 'byte'), | |
| } | |
| def __init__(self, *args, **kwargs): | |
| super(PackedSequenceTest, self).__init__(*args, **kwargs) | |
| self.batch_size = 5 | |
| self.max_length = 6 | |
| def _ordered_sequence(self, tensor_type): | |
| """Create ordered list of random sequences""" | |
| seqs = [tensor_type(random.randint(1, self.max_length)) | |
| for _ in range(self.batch_size)] | |
| if tensor_type == torch.ByteTensor: | |
| seqs = [s.random_(0, 256) for s in seqs] | |
| else: | |
| seqs = [s.random_(-128, 128) for s in seqs] | |
| ordered = sorted(seqs, key=len, reverse=True) | |
| return ordered | |
| def _padded_sequence(self, tensor_type): | |
| """Create Tensor of random padded sequences""" | |
| ordered = self._ordered_sequence(tensor_type) | |
| lengths = [len(i) for i in ordered] | |
| padded_tensor = rnn_utils.pad_sequence(ordered) | |
| return padded_tensor, lengths | |
| def test_type_casts(self): | |
| """Test type casting of `PackedSequence` against type casting of tensor""" | |
| for _, (input_type, _) in self._type_by_name.items(): | |
| for expected_type_str, (_, cast_str) in self._type_by_name.items(): | |
| for enforce_sorted in [True, False]: | |
| padded, lengths = self._padded_sequence(input_type) | |
| packed = rnn_utils.pack_padded_sequence( | |
| padded, lengths, enforce_sorted=enforce_sorted) | |
| # Apply cast to `PackedSequence` instance and unpack | |
| masked = getattr(packed, cast_str)() | |
| unpacked, lengths_out = rnn_utils.pad_packed_sequence(masked) | |
| self.assertEqual(unpacked.type(), expected_type_str) | |
| def test_wrong_order(self): | |
| a = torch.ones(25, 300) | |
| b = torch.ones(22, 300) | |
| b_a = rnn_utils.pad_sequence([b, a]) | |
| self.assertRaises( | |
| RuntimeError, | |
| lambda: rnn_utils.pack_padded_sequence(b_a, [22, 25], enforce_sorted=True)) | |
| def test_pad_sequence_with_tensor_sequences(self): | |
| seq_tuple_input = torch.nn.utils.rnn.pad_sequence( | |
| (torch.tensor([[7, 6]]), torch.tensor([[-7, -1]])) | |
| ) | |
| seq_tensor_input = torch.nn.utils.rnn.pad_sequence( | |
| torch.tensor([[[7, 6]], [[-7, -1]]]) | |
| ) | |
| self.assertEqual(seq_tuple_input, seq_tensor_input) | |
| self.assertEqual(seq_tuple_input.shape, torch.Size([1, 2, 2])) | |
| def test_pad_sequence_with_non_iterable_sequences(self): | |
| msg = r"Expected iterable for input sequences, but got arg of type" | |
| with self.assertRaisesRegex(RuntimeError, msg): | |
| torch.nn.utils.rnn.pad_sequence(5) | |
| def test_total_length(self): | |
| padded, lengths = self._padded_sequence(torch.FloatTensor) | |
| max_length = max(lengths) | |
| packed = rnn_utils.pack_padded_sequence(padded, lengths) | |
| # test ValueError if total_length < max_length | |
| for total_length in (-1, 0, max_length - 1): | |
| for batch_first in (True, False): | |
| def err_fn(): | |
| rnn_utils.pad_packed_sequence(packed, batch_first=batch_first, | |
| total_length=total_length) | |
| self.assertRaisesRegex(ValueError, | |
| r'Expected total_length to be at least the ' | |
| r'length of the longest sequence in input', | |
| err_fn) | |
| # test that pad_packed_sequence returns results of correct length | |
| for batch_first in (True, False): | |
| no_extra_pad, _ = rnn_utils.pad_packed_sequence(packed, batch_first=batch_first) | |
| for total_length_delta in (0, 1, 8): | |
| total_length = max_length + total_length_delta | |
| unpacked, lengths_out = rnn_utils.pad_packed_sequence(packed, batch_first=batch_first, | |
| total_length=total_length) | |
| self.assertEqual(lengths, lengths_out) | |
| self.assertEqual(unpacked.size(1 if batch_first else 0), total_length) | |
| if total_length_delta == 0: | |
| ref_output = no_extra_pad | |
| elif batch_first: | |
| extra_pad = no_extra_pad.new_zeros(self.batch_size, total_length_delta) | |
| ref_output = torch.cat([no_extra_pad, extra_pad], 1) | |
| else: | |
| extra_pad = no_extra_pad.new_zeros(total_length_delta, self.batch_size) | |
| ref_output = torch.cat([no_extra_pad, extra_pad], 0) | |
| self.assertEqual(unpacked, ref_output) | |
| def test_to(self): | |
| for enforce_sorted in (True, False): | |
| padded, lengths = self._padded_sequence(torch.IntTensor) | |
| a = rnn_utils.pack_padded_sequence( | |
| padded, lengths, enforce_sorted=enforce_sorted).cpu() | |
| self.assertIs(a, a.to('cpu')) | |
| self.assertIs(a, a.cpu()) | |
| self.assertIs(a, a.to('cpu', dtype=torch.int32)) | |
| self.assertEqual(a.long(), a.to(torch.int64)) | |
| if torch.cuda.is_available(): | |
| for cuda in ['cuda', 'cuda:0' if torch.cuda.device_count() == 1 else 'cuda:1']: | |
| b = a.cuda(device=cuda) | |
| self.assertIs(b, b.to(cuda)) | |
| self.assertIs(b, b.cuda()) | |
| self.assertEqual(a, b.to('cpu')) | |
| self.assertEqual(b, a.to(cuda)) | |
| self.assertEqual(a, b.to('cpu', dtype=torch.int32)) | |
| self.assertIs(b, b.to(dtype=torch.int32)) | |
| self.assertEqual(b.long(), b.to(dtype=torch.int64)) | |
| def test_to_memory_format(self): | |
| m = torch.nn.Conv2d(in_channels=16, out_channels=32, kernel_size=2, bias=True) | |
| m = m.to(memory_format=torch.channels_last) | |
| for param in m.parameters(): | |
| if param.dim() == 4: | |
| self.assertTrue(param.is_contiguous(memory_format=torch.channels_last)) | |
| def test_pad_sequence(self): | |
| def pad(tensor, length): | |
| return torch.cat( | |
| [tensor.data, tensor.data.new( | |
| length - tensor.size(0), *tensor.size()[1:]).zero_()]) | |
| # single dimensional | |
| a = torch.tensor([1, 2, 3]) | |
| b = torch.tensor([4, 5]) | |
| c = torch.tensor([6]) | |
| # batch_first = true | |
| expected = torch.tensor([[4, 5, 0], [1, 2, 3], [6, 0, 0]]) | |
| padded = rnn_utils.pad_sequence([b, a, c], True) | |
| self.assertEqual(padded, expected) | |
| # batch_first = false | |
| padded = rnn_utils.pad_sequence([b, a, c]) | |
| self.assertEqual(padded, expected.transpose(0, 1)) | |
| # pad with non-zero value | |
| expected = torch.tensor([[4, 5, 1], [1, 2, 3], [6, 1, 1]]) | |
| padded = rnn_utils.pad_sequence([b, a, c], True, 1) | |
| self.assertEqual(padded, expected) | |
| # Test pad sorted sequence | |
| expected = torch.tensor([[1, 2, 3], [4, 5, 0], [6, 0, 0]]) | |
| padded = rnn_utils.pad_sequence([a, b, c], True) | |
| self.assertEqual(padded, expected) | |
| # more dimensions | |
| maxlen = 9 | |
| for num_dim in (0, 1, 2, 3): | |
| sequences = [] | |
| trailing_dims = [4] * num_dim | |
| for i in range(1, maxlen + 1): | |
| seq_len = i * i | |
| sequences.append(torch.rand(seq_len, 5, *trailing_dims)) | |
| random.shuffle(sequences) | |
| expected = [] | |
| for seq in sequences: | |
| expected.append(pad(seq, maxlen * maxlen)) | |
| # batch first = true | |
| expected = torch.stack(expected) | |
| padded = rnn_utils.pad_sequence(sequences, True) | |
| self.assertEqual(padded, expected) | |
| # batch first = false | |
| padded = rnn_utils.pad_sequence(sequences) | |
| self.assertEqual(padded, expected.transpose(0, 1)) | |
| def test_unpad_sequence(self): | |
| # single dimensional | |
| a = torch.tensor([1, 2, 3]) | |
| b = torch.tensor([4, 5]) | |
| c = torch.tensor([6]) | |
| sequences = [a, b, c] | |
| lengths = torch.as_tensor([v.size(0) for v in sequences]) | |
| for batch_first in [True, False]: | |
| padded_sequences = rnn_utils.pad_sequence(sequences, batch_first=batch_first) | |
| unpadded_sequences = rnn_utils.unpad_sequence(padded_sequences, lengths, batch_first=batch_first) | |
| self.assertEqual(sequences, unpadded_sequences) | |
| # more dimensions | |
| maxlen = 9 | |
| for num_dim in (0, 1, 2, 3): | |
| sequences = [] | |
| trailing_dims = [4] * num_dim | |
| for i in range(1, maxlen + 1): | |
| seq_len = i * i | |
| sequences.append(torch.rand(seq_len, 5, *trailing_dims)) | |
| random.shuffle(sequences) | |
| lengths = torch.as_tensor([v.size(0) for v in sequences]) | |
| padded_sequences = rnn_utils.pad_sequence(sequences, batch_first=batch_first) | |
| unpadded_sequences = rnn_utils.unpad_sequence(padded_sequences, lengths, batch_first=batch_first) | |
| self.assertEqual(sequences, unpadded_sequences) | |
| def test_pack_sequence(self): | |
| def _compatibility_test(sequences, lengths, batch_first, enforce_sorted=False): | |
| padded = rnn_utils.pad_sequence(sequences, batch_first) | |
| packed = rnn_utils.pack_sequence(sequences, enforce_sorted) | |
| unpacked = rnn_utils.pad_packed_sequence(packed, batch_first) | |
| self.assertEqual(padded, unpacked[0]) | |
| pack_padded = rnn_utils.pack_padded_sequence( | |
| padded, lengths, batch_first, enforce_sorted) | |
| self.assertEqual(packed, pack_padded) | |
| # single dimensional | |
| a = torch.tensor([1, 2, 3]) | |
| b = torch.tensor([4, 5]) | |
| c = torch.tensor([6]) | |
| packed = rnn_utils.pack_sequence([a, b, c], enforce_sorted=False) | |
| expected = torch.tensor([1, 4, 6, 2, 5, 3]) | |
| self.assertEqual(packed.batch_sizes, [3, 2, 1]) | |
| self.assertEqual(packed.data.data, expected) | |
| self.assertEqual(packed.sorted_indices, [0, 1, 2]) | |
| self.assertEqual(packed.unsorted_indices, [0, 1, 2]) | |
| packed_unsorted = rnn_utils.pack_sequence([b, c, a], enforce_sorted=False) | |
| self.assertEqual(packed_unsorted.batch_sizes, [3, 2, 1]) | |
| self.assertEqual(packed_unsorted.data.data, expected) | |
| self.assertEqual(packed_unsorted.sorted_indices, [2, 0, 1]) | |
| self.assertEqual(packed_unsorted.unsorted_indices, [1, 2, 0]) | |
| # single dimensional, enforce_sorted = True | |
| packed_enforce_sorted = rnn_utils.pack_sequence([a, b, c], enforce_sorted=True) | |
| self.assertEqual(packed_enforce_sorted.batch_sizes, [3, 2, 1]) | |
| self.assertEqual(packed_enforce_sorted.data.data, expected) | |
| self.assertTrue(packed_enforce_sorted.sorted_indices is None) | |
| self.assertTrue(packed_enforce_sorted.unsorted_indices is None) | |
| with self.assertRaisesRegex(RuntimeError, 'must be sorted in decreasing order'): | |
| rnn_utils.pack_sequence([b, c, a], enforce_sorted=True) | |
| with self.assertRaisesRegex(RuntimeError, 'You can pass `enforce_sorted=False`'): | |
| rnn_utils.pack_sequence([b, c, a], enforce_sorted=True) | |
| # more dimensions | |
| maxlen = 9 | |
| for num_dim in (0, 1, 2, 3): | |
| sequences = [] | |
| lengths = [] | |
| trailing_dims = [4] * num_dim | |
| for i in range(maxlen, 0, -1): | |
| seq_len = i * i | |
| lengths.append(seq_len) | |
| sequences.append(torch.rand(seq_len, 5, *trailing_dims)) | |
| unsorted_sequences = [s.clone() for s in sequences] | |
| random.shuffle(unsorted_sequences) | |
| unsorted_sequences_lengths = [t.size(0) for t in unsorted_sequences] | |
| # compatibility with other utilities | |
| for batch_first in (True, False): | |
| for enforce_sorted in (True, False): | |
| _compatibility_test(sequences, lengths, batch_first, enforce_sorted) | |
| _compatibility_test(unsorted_sequences, unsorted_sequences_lengths, | |
| batch_first) | |
| def test_unpack_sequence(self): | |
| # single dimensional | |
| a = torch.tensor([1, 2, 3]) | |
| b = torch.tensor([4, 5]) | |
| c = torch.tensor([6]) | |
| sequences = [a, b, c] | |
| packed_sequences = rnn_utils.pack_sequence(sequences, enforce_sorted=False) | |
| unpacked_sequences = rnn_utils.unpack_sequence(packed_sequences) | |
| self.assertEqual(sequences, unpacked_sequences) | |
| # more dimensions | |
| maxlen = 9 | |
| for num_dim in (0, 1, 2, 3): | |
| sequences = [] | |
| trailing_dims = [4] * num_dim | |
| for i in range(1, maxlen + 1): | |
| seq_len = i * i | |
| sequences.append(torch.rand(seq_len, 5, *trailing_dims)) | |
| random.shuffle(sequences) | |
| packed_sequences = rnn_utils.pack_sequence(sequences, enforce_sorted=False) | |
| unpacked_sequences = rnn_utils.unpack_sequence(packed_sequences) | |
| self.assertEqual(sequences, unpacked_sequences) | |
| def test_pack_padded_sequence(self): | |
| def generate_test_case(sorted_lengths, should_shuffle): | |
| def pad(tensor, length): | |
| return torch.cat([tensor, tensor.new(length - tensor.size(0), *tensor.size()[1:]).zero_()]) | |
| max_length = sorted_lengths[0] | |
| batch_sizes = [sum(map(bool, filter(lambda x: x >= i, sorted_lengths))) | |
| for i in range(1, max_length + 1)] | |
| offset = 0 | |
| padded = torch.cat([pad(i * 100 + torch.arange(1., 5 * l + 1).view(l, 1, 5), max_length) | |
| for i, l in enumerate(sorted_lengths, 1)], 1) | |
| expected_data = [[torch.arange(1., 6) + (i + 1) * 100 + 5 * n for i in range(batch_size)] | |
| for n, batch_size in enumerate(batch_sizes)] | |
| expected_data = list(itertools.chain.from_iterable(expected_data)) | |
| expected_data = torch.stack(expected_data, dim=0) | |
| if should_shuffle: | |
| # Shuffle the padded sequence to create an unsorted sequence | |
| permutation = list(range(len(sorted_lengths))) | |
| random.shuffle(permutation) | |
| unsorted_indices = torch.tensor(permutation) | |
| padded = padded.index_select(1, unsorted_indices) | |
| lengths = torch.tensor(sorted_lengths).index_select(0, unsorted_indices) | |
| else: | |
| unsorted_indices = None | |
| lengths = sorted_lengths | |
| return padded.requires_grad_(), lengths, expected_data, batch_sizes, unsorted_indices | |
| test_cases = [ | |
| # sorted_lengths, should_shuffle | |
| [[10, 8, 4, 2, 2, 2, 1], False], | |
| [[11, 10, 8, 6, 4, 3, 1], False], | |
| [[11, 10, 8, 6, 4, 3, 1], True], | |
| ] | |
| for test_case, batch_first in itertools.product(test_cases, (True, False)): | |
| sorted_lengths, should_shuffle = test_case | |
| padded, lengths, expected_data, batch_sizes, unsorted_indices = generate_test_case( | |
| sorted_lengths, should_shuffle) | |
| src = padded | |
| if batch_first: | |
| src = src.transpose(0, 1) | |
| # check output | |
| packed = rnn_utils.pack_padded_sequence(src, lengths, batch_first=batch_first, | |
| enforce_sorted=not should_shuffle) | |
| self.assertEqual(packed.data.data, expected_data) | |
| self.assertEqual(packed.batch_sizes, batch_sizes) | |
| self.assertEqual(packed.unsorted_indices, unsorted_indices) | |
| # test inverse | |
| unpacked, unpacked_len = rnn_utils.pad_packed_sequence(packed, batch_first=batch_first) | |
| self.assertEqual(unpacked, src) | |
| self.assertEqual(unpacked_len, lengths) | |
| # check grad | |
| if padded.grad is not None: | |
| padded.grad.data.zero_() | |
| grad_output = unpacked.data.clone().normal_() | |
| unpacked.backward(grad_output) | |
| if batch_first: | |
| grad_output.transpose_(0, 1) | |
| for i, l in enumerate(lengths): | |
| self.assertEqual(padded.grad.data[:l, i], grad_output[:l, i]) | |
| if l < 10: | |
| self.assertEqual(padded.grad.data[l:, i].abs().sum(), 0) | |
| # test error messages | |
| with self.assertRaisesRegex(RuntimeError, 'You can pass `enforce_sorted=False`'): | |
| packed = rnn_utils.pack_padded_sequence(torch.randn(3, 3), [1, 3, 2]) | |
| with self.assertRaisesRegex(RuntimeError, 'empty tensor'): | |
| packed = rnn_utils.pack_padded_sequence(torch.randn(0, 0), []) | |
| if __name__ == '__main__': | |
| run_tests() |