| # Owner(s): ["module: nn"] | |
| import unittest | |
| import random | |
| import itertools | |
| from itertools import product | |
| import torch | |
| from torch.testing._internal.common_utils import run_tests, set_default_dtype, \ | |
| instantiate_parametrized_tests, parametrize as parametrize_test, _assertGradAndGradgradChecks | |
| from torch.testing._internal.common_cuda import TEST_CUDA | |
| from torch.testing._internal.common_nn import NNTestCase | |
| from torch.testing._internal.common_device_type import onlyNativeDeviceTypes, dtypes, \ | |
| instantiate_device_type_tests, dtypesIfCUDA, onlyCUDA, \ | |
| TEST_WITH_ROCM, skipCUDAIf, skipMeta | |
| import torch.nn.functional as F | |
| import torch.nn as nn | |
| from torch.testing._internal.common_utils import dtype2prec_DONTUSE | |
| class TestEmbeddingNN(NNTestCase): | |
| _do_cuda_memory_leak_check = True | |
| _do_cuda_non_default_stream = True | |
| @unittest.skipIf(not TEST_CUDA, "CUDA unavailable") | |
| def test_embedding_max_norm_unsorted_repeating_indices(self): | |
| def create_embedding(device): | |
| # Seed RNG so we get the same Embedding each time | |
| torch.manual_seed(0) | |
| return torch.nn.Embedding( | |
| num_embeddings=20, | |
| embedding_dim=64, | |
| max_norm=1.0).to(device) | |
| ix = torch.arange(2, device='cpu', dtype=torch.long).repeat(2000) | |
| out_cpu = create_embedding('cpu')(ix) | |
| ix = ix.to('cuda') | |
| out = create_embedding('cuda')(ix) | |
| self.assertEqual(out.cpu(), out_cpu) | |
| def test_embedding_sparse_basic(self): | |
| embedding = nn.Embedding(10, 20, sparse=True) | |
| input = torch.tensor([[0, 2, 4, 5], [4, 3, 0, 9]], dtype=torch.long) | |
| embedding(input).sum().backward() | |
| self.assertTrue(embedding.weight.grad.is_sparse) | |
| self.assertEqual(embedding.weight.grad.shape, embedding.weight.shape) | |
| def test_embedding_sparse_empty_tensor(self): | |
| embedding = nn.Embedding(0, 0, sparse=True) | |
| input = torch.tensor([], dtype=torch.int64) | |
| embedding(input).sum().backward() | |
| self.assertTrue(embedding.weight.grad.is_sparse) | |
| self.assertEqual(embedding.weight.grad.shape, embedding.weight.shape) | |
| embedding = nn.Embedding(10, 0, sparse=True) | |
| input = torch.LongTensor([[0, 2, 4, 5], [4, 3, 0, 9]]) | |
| embedding(input).sum().backward() | |
| self.assertTrue(embedding.weight.grad.is_sparse) | |
| self.assertEqual(embedding.weight.grad.shape, embedding.weight.shape) | |
| def test_move_sparse_half_embedding(self): | |
| embedding = nn.Embedding(10, 3, sparse=True) | |
| self.assertEqual(embedding.weight.device.type, 'cpu') | |
| self.assertEqual(embedding.weight.dtype, torch.get_default_dtype()) | |
| embedding.to(torch.float16) | |
| self.assertEqual(embedding.weight.dtype, torch.float16) | |
| self.assertEqual(embedding.embedding_dim, 3) | |
| self.assertEqual(embedding.num_embeddings, 10) | |
| if torch.cuda.is_available(): | |
| embedding.to('cuda') | |
| self.assertEqual(embedding.weight.device.type, 'cuda') | |
| embedding.to('cpu') | |
| self.assertEqual(embedding.weight.device.type, 'cpu') | |
| def test_embedding_max_norm(self): | |
| embedding = nn.Embedding(22, 5, max_norm=1.0) | |
| input = torch.tensor([2, 8, 8, 6], dtype=torch.long) | |
| output = embedding(input) | |
| self.assertEqual(output[1], output[2]) | |
| self.assertTrue(output.data.norm(p=2, dim=1).le(1).all()) | |
| @parametrize_test("dtype", (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64, torch.float, torch.double)) | |
| def test_embedding_from_pretrained(self, dtype): | |
| a = torch.tensor([[1., 2., 3.], [4., 5., 6.]], dtype=dtype) | |
| embedding = nn.Embedding.from_pretrained(a) | |
| self.assertEqual(a, embedding.weight.data) | |
| input = torch.LongTensor([0, 1]) | |
| output = embedding(input) | |
| self.assertEqual(a, output) | |
| def test_embedding_bag_from_pretrained(self): | |
| a = torch.tensor([[1., 2., 3.], [4., 5., 6.]]) | |
| embedding = nn.EmbeddingBag.from_pretrained(a) | |
| self.assertEqual(a, embedding.weight) | |
| input = torch.tensor([0, 1], dtype=torch.long) | |
| output = embedding(input, torch.arange(input.size(0))) | |
| self.assertEqual(a, output) | |
| def test_embedding_from_pretrained_padding_idx(self): | |
| padding_idx = 2 | |
| padding_vec = torch.ones(3) * 7 | |
| embeddings = torch.rand(4, 3, requires_grad=True) | |
| with torch.no_grad(): | |
| embeddings[padding_idx] = padding_vec | |
| embedding_nn = nn.Embedding.from_pretrained(embeddings, padding_idx=padding_idx) | |
| self.assertEqual(embedding_nn.weight[padding_idx], padding_vec) | |
| def test_embedding_bag_from_pretrained_padding_idx(self): | |
| padding_idx = 2 | |
| embeddings = torch.rand(4, 3, requires_grad=True) | |
| embedding_nn = nn.EmbeddingBag.from_pretrained(embeddings, padding_idx=padding_idx) | |
| self.assertEqual(embedding_nn.weight, embeddings) | |
| def test_embedding_from_pretrained_options(self): | |
| with set_default_dtype(torch.double): | |
| a = torch.tensor([[1., 2., 3.], [4., 5., 6.]]) | |
| opts = { | |
| "max_norm": 2., | |
| "norm_type": .5, | |
| "scale_grad_by_freq": False, | |
| "sparse": True | |
| } | |
| embedding = nn.Embedding.from_pretrained(a, **opts) | |
| input = torch.LongTensor([0, 1]) | |
| output = embedding(input) | |
| # test output and that weight matrix was renormalized | |
| self.assertEqual(a, output) | |
| self.assertTrue(a.ne(torch.arange(1, 7, dtype=a.dtype).view(2, 3)).all()) | |
| self.assertTrue(output.data.norm(p=opts["norm_type"], dim=1).le(opts["max_norm"]).all()) | |
| def test_embedding_functional(self): | |
| a = torch.tensor([ | |
| [1, 3, 2], | |
| [0, 2, 1] | |
| ], dtype=torch.long) | |
| embeddings = torch.rand(4, 3, requires_grad=True) | |
| embed_old = torch.nn.Embedding(4, 3) | |
| embed_old.weight.data = embeddings.data | |
| res_old = embed_old(a) | |
| res_F = F.embedding(a, embeddings) | |
| self.assertEqual(res_old, res_F) | |
| embed_old = torch.nn.Embedding(4, 3) | |
| embed_old = embed_old.from_pretrained(embeddings, padding_idx=2) | |
| res_old = embed_old(a) | |
| res_F = F.embedding(a, embeddings, padding_idx=2) | |
| self.assertEqual(res_old, res_F) | |
| def test_embedding_bag_functional(self): | |
| a = torch.tensor([ | |
| [1, 3, 2], | |
| [0, 2, 1] | |
| ], dtype=torch.long) | |
| embeddings = torch.rand(4, 3, requires_grad=True) | |
| embed_old = torch.nn.EmbeddingBag(4, 3) | |
| embed_old.weight = torch.nn.Parameter(embeddings) | |
| res_old = embed_old(a) | |
| res_F = F.embedding_bag(a, embeddings) | |
| self.assertEqual(res_old, res_F) | |
| embed_old = torch.nn.EmbeddingBag(4, 3) | |
| embed_old = embed_old.from_pretrained(embeddings, padding_idx=2) | |
| res_old = embed_old(a) | |
| res_F = F.embedding_bag(a, embeddings, padding_idx=2) | |
| self.assertEqual(res_old, res_F) | |
| # Make sure that error is thrown if padding_idx is out of bounds | |
| def test_embedding_bag_padding_idx_error(self): | |
| a = torch.tensor([ | |
| [1, 3, 2], | |
| [0, 2, 1] | |
| ], dtype=torch.long) | |
| num_embeddings = 4 | |
| num_features = 3 | |
| embeddings = torch.rand(num_embeddings, num_features, requires_grad=True) | |
| functional_err_msg = r'padding_idx must be within the number of embeddings' | |
| module_err_msg = r'padding_idx must be within num_embeddings' | |
| for padding_idx in range(-(num_embeddings + 2), (num_embeddings + 2)): | |
| if (padding_idx < -num_embeddings) or (padding_idx >= num_embeddings): | |
| with self.assertRaisesRegex(RuntimeError, functional_err_msg): | |
| F.embedding_bag(a, embeddings, padding_idx=padding_idx) | |
| with self.assertRaisesRegex(AssertionError, module_err_msg): | |
| torch.nn.EmbeddingBag(num_embeddings, num_features, padding_idx=padding_idx) | |
| else: | |
| F.embedding_bag(a, embeddings, padding_idx=padding_idx) | |
| torch.nn.EmbeddingBag(num_embeddings, num_features, padding_idx=padding_idx) | |
| def test_embeddingbag_from_pretrained(self): | |
| a = torch.tensor([[1., 2., 3.], [4., 5., 6.]]) | |
| embeddingbag = nn.EmbeddingBag.from_pretrained(a) | |
| self.assertEqual(a, embeddingbag.weight.data) | |
| input = torch.LongTensor([[0, 1]]) | |
| output = embeddingbag(input) | |
| self.assertEqual(a.mean(0, keepdim=True), output) | |
| def test_embeddingbag_from_pretrained_options(self): | |
| with set_default_dtype(torch.double): | |
| a = torch.tensor([[1., 2., 3.], [4., 5., 6.]]) | |
| opts = { | |
| "max_norm": 2., | |
| "norm_type": .5, | |
| "scale_grad_by_freq": False, | |
| "mode": "max", | |
| "sparse": False | |
| } | |
| embeddingbag = nn.EmbeddingBag.from_pretrained(a, **opts) | |
| input = torch.LongTensor([[0, 1]]) | |
| output = embeddingbag(input) | |
| self.assertEqual(a.max(0, keepdim=True)[0], output) | |
| self.assertTrue(a.ne(torch.arange(1, 7, dtype=a.dtype).view(2, 3)).all()) | |
| self.assertTrue(a.norm(p=opts["norm_type"], dim=1).le(opts["max_norm"]).all()) | |
| class TestEmbeddingNNDeviceType(NNTestCase): | |
| def test_embedding_dense_grad(self, device): | |
| with set_default_dtype(torch.double): | |
| embd = nn.Embedding(20, 20).to(device) | |
| weight = embd.weight | |
| def fn_wrapper(device): | |
| def fn(weight): | |
| inp = torch.tensor([[0, 1, 1, 2], [3, 5, 7, 11]], dtype=torch.long).to(device) | |
| return torch.nn.functional.embedding(inp, weight) | |
| return fn | |
| fn = fn_wrapper(device) | |
| _assertGradAndGradgradChecks(self, fn, (weight, )) | |
| def test_embedding_scalar_weight_error(self, device): | |
| indices = torch.rand(2, 2, device=device).long() | |
| weights = [ | |
| torch.tensor(1.0, device=device), | |
| torch.tensor(1.0, device=device).reshape(1, 1, 1), | |
| ] | |
| for weight in weights: | |
| with self.assertRaisesRegex(RuntimeError, "'weight' must be 2-D"): | |
| torch.nn.functional.embedding(indices, weight) | |
| @dtypesIfCUDA(torch.float16, torch.float64) | |
| @dtypes(torch.float64) | |
| def test_embedding_backward(self, device, dtype): | |
| embedding = nn.Embedding(10, 3, sparse=True) | |
| tensor = torch.tensor([[7, 1, 3]]) | |
| ones = torch.tensor(1., dtype=dtype).expand(3, 3) | |
| tensorTwice = tensor.repeat(1, 2) | |
| onesTwice = torch.cat((ones, ones)) | |
| embedding = embedding.to(dtype=dtype).to(device) | |
| tensor = tensor.to(device) | |
| ones = ones.to(device) | |
| tensorTwice = tensorTwice.to(device) | |
| onesTwice = onesTwice.to(device) | |
| embedding.zero_grad() | |
| embedding(tensor[0]).sum().backward() | |
| self.assertEqual(embedding.weight.grad._indices(), tensor) | |
| self.assertEqual(embedding.weight.grad._values(), ones) | |
| embedding.zero_grad() | |
| embedding(tensor[0]).sum().backward() | |
| embedding(tensor[0]).sum().backward() | |
| self.assertEqual(embedding.weight.grad._indices(), tensorTwice) | |
| self.assertEqual(embedding.weight.grad._values(), onesTwice) | |
| embedding.zero_grad() | |
| embedding(tensor[0]).sum().backward() | |
| tensor[0, 0] = 8 | |
| embedding(tensor[0]).sum().backward() | |
| tensorTwice[0, 3] = 8 | |
| self.assertEqual(embedding.weight.grad._indices(), tensorTwice) | |
| self.assertEqual(embedding.weight.grad._values(), onesTwice) | |
| @dtypesIfCUDA(*((torch.float, torch.double, torch.bfloat16, torch.half) | |
| if TEST_WITH_ROCM else (torch.float, torch.double, torch.half))) | |
| @dtypes(torch.float32) | |
| def test_embedding_max_norm_backward(self, device, dtype): | |
| # can't use gradcheck since in place renorm makes analytical gradients different from produced ones | |
| weight = torch.randn((4, 4), device=device, dtype=dtype) * 2 | |
| weight.requires_grad_() | |
| inp_list = [0, 1, 2, 2] | |
| inp = torch.tensor(inp_list, device=device) | |
| out = nn.functional.embedding(inp, weight, max_norm=1.).sum() | |
| out.backward() | |
| expected_grad = torch.tensor([[1., 1., 2., 0.]], device=device, dtype=dtype).transpose(0, 1).expand(4, 4) | |
| self.assertEqual(weight.grad, expected_grad) | |
| @dtypesIfCUDA(*((torch.float, torch.double, torch.bfloat16, torch.half) | |
| if TEST_WITH_ROCM else (torch.float, torch.double, torch.half))) | |
| @dtypes(torch.float32) | |
| def test_embedding_max_norm_fwd_AD(self, device, dtype): | |
| if torch.device(device).type == 'xla': | |
| self.skipTest("forward AD doesn't work on xla") | |
| # can't use gradcheck since in place renorm makes analytical gradients different from produced ones | |
| weight = torch.randn((4, 4), device=device, dtype=dtype) * 2 | |
| tangent = torch.ones((4, 4), device=device, dtype=dtype) | |
| inp = torch.tensor([[0, 1], [2, 2]], device=device) | |
| with torch.autograd.forward_ad.dual_level(): | |
| dual_weight = torch.autograd.forward_ad.make_dual(weight, tangent) | |
| out = nn.functional.embedding(inp, dual_weight, max_norm=1.) | |
| jvp = torch.autograd.forward_ad.unpack_dual(out).tangent | |
| expected_grad = torch.ones((2, 2, 4), device=device, dtype=dtype) | |
| self.assertEqual(jvp, expected_grad) | |
| @dtypesIfCUDA(*((torch.float, torch.double, torch.bfloat16, torch.half) | |
| if TEST_WITH_ROCM else (torch.float, torch.double, torch.half))) | |
| @dtypes(torch.float32) | |
| def test_embedding_padding_idx(self, device, dtype): | |
| embedding = nn.Embedding(10, 20, padding_idx=0).to(device, dtype) | |
| input = torch.tensor([[0, 2, 4, 5], [4, 3, 0, 9]], dtype=torch.long).to(device) | |
| output = embedding(input) | |
| self.assertEqual(output[0][0].sum(), 0) | |
| self.assertEqual(output[1][2].sum(), 0) | |
| embedding = nn.Embedding(10, 20, padding_idx=0, sparse=True).to(device, dtype) | |
| input = torch.tensor([[0, 2, 4, 5], [4, 3, 0, 9]], dtype=torch.long).to(device) | |
| output = embedding(input) | |
| self.assertEqual(output[0][0].sum(), 0) | |
| self.assertEqual(output[1][2].sum(), 0) | |
| # negative indexing check for padding_idx | |
| # padding_idx=-2, num_embeddings=10 ==> index 8 padded | |
| embedding = nn.Embedding(10, 20, padding_idx=-2).to(device, dtype) | |
| input = torch.tensor([[0, 2, 8, 5], [4, 8, 0, 9]], dtype=torch.long).to(device) | |
| output = embedding(input) | |
| self.assertEqual(output[0][2].sum(), 0) | |
| self.assertEqual(output[1][1].sum(), 0) | |
| embedding = nn.Embedding(10, 20, padding_idx=-2, sparse=True).to(device, dtype) | |
| input = torch.tensor([[0, 2, 8, 5], [4, 8, 0, 9]], dtype=torch.long).to(device) | |
| output = embedding(input) | |
| self.assertEqual(output[0][2].sum(), 0) | |
| self.assertEqual(output[1][1].sum(), 0) | |
| # change padding vector | |
| padding_vector = torch.ones(20, dtype=dtype, device=device) | |
| embedding = nn.Embedding(10, 20, padding_idx=2, sparse=True).to(device, dtype) | |
| with torch.no_grad(): | |
| embedding.weight[2] = padding_vector | |
| input = torch.tensor([0, 2], dtype=torch.long).to(device) | |
| output = embedding(input) | |
| self.assertEqual(output[1], padding_vector) | |
| # out of bounds check for padding_idx | |
| self.assertRaises(AssertionError, nn.Embedding, num_embeddings=10, embedding_dim=20, padding_idx=25) | |
| self.assertRaises(AssertionError, nn.Embedding, num_embeddings=10, embedding_dim=20, padding_idx=-25) | |
| padding_idx = 0 | |
| embedding = nn.Embedding(5, 2, padding_idx=padding_idx).to(device, dtype) | |
| for n in (1, 2, 1000): # Need large N to trigger all the methods we have implemented | |
| for other_indices in ([], [1, 3], [2]): | |
| indices = torch.tensor(other_indices + [padding_idx] * n, dtype=torch.long).to(device) | |
| pre = embedding.weight[padding_idx].clone() | |
| embedding(indices).sum().backward() | |
| after = (embedding.weight + embedding.weight.grad)[padding_idx] | |
| embedding.zero_grad() | |
| self.assertEqual(after, pre) | |
| # test double backward | |
| emb_sum = embedding(indices).sum() | |
| emb_grad = torch.autograd.grad(outputs=emb_sum, inputs=list(embedding.parameters()), retain_graph=True) | |
| scalar = emb_grad[0].sum() + emb_sum | |
| scalar.backward() | |
| after = (embedding.weight + embedding.weight.grad)[padding_idx] | |
| embedding.zero_grad() | |
| self.assertEqual(after, pre) | |
| # Check correctness of torch.nn.functional.embedding_bag forward and | |
| # backward functions with padding_idx, given a 1D input separated into bags | |
| # with an offset array. Compare against an equivalent 2D input that uses | |
| # padding indices to fill in the gaps indicated by the offset array | |
| @onlyNativeDeviceTypes | |
| @dtypes(torch.float32, torch.float64) | |
| @dtypesIfCUDA(torch.half, torch.bfloat16) | |
| def test_embedding_bag_1D_padding_idx(self, device, dtype): | |
| num_features = 3 | |
| max_indices_per_bag = 10 | |
| num_bags = 10 | |
| num_words = 100 | |
| def gen_1D_indices_offsets(include_last_offset, allpad): | |
| indices = [] | |
| offsets = [] | |
| cur_offset = 0 | |
| # Make one bag full and one bag empty, for extra coverage | |
| empty_bag = random.randint(0, num_bags - 1) | |
| full_bag = empty_bag | |
| while full_bag == empty_bag: | |
| full_bag = random.randint(0, num_bags - 1) | |
| for bag in range(num_bags): | |
| offsets.append(cur_offset) | |
| if bag == full_bag: | |
| bag_size = max_indices_per_bag | |
| elif bag == empty_bag: | |
| bag_size = 0 | |
| else: | |
| bag_size = random.randint(1, max_indices_per_bag - 1) | |
| indices += [1 if allpad else random.randint(0, num_words - 1) for _ in range(bag_size)] | |
| cur_offset += bag_size | |
| # embedding_bag requires first entry of offsets to be 0 | |
| assert offsets[0] == 0 | |
| indices = torch.tensor(indices, device=device) | |
| if include_last_offset: | |
| offsets.append(indices.size(0)) | |
| offsets = torch.tensor(offsets, device=device) | |
| return indices, offsets | |
| # Convert a 1-D indices-offsets representation into 2-D. Fill any empty | |
| # indices with padding_idx | |
| def gen_2D_indices_from_1D(indices_1D, offsets, include_last_offset, padding_idx): | |
| assert offsets[0] == 0 | |
| if include_last_offset: | |
| offsets = offsets[:-1] | |
| indices_2D = torch.empty(num_bags, max_indices_per_bag, device=device, dtype=torch.long) | |
| for bag in range(num_bags): | |
| # Determine the start and end position of the bag within indices_1D | |
| start = offsets[bag] | |
| end = len(indices_1D) if bag + 1 == num_bags else offsets[bag + 1] | |
| end = min(len(indices_1D), end) | |
| # Pull out the bag's indices from indices_1D, and fill any | |
| # remaining space with padding indices | |
| indices_in_bag = [] | |
| for item_pos in range(0, max_indices_per_bag): | |
| if (start + item_pos) < end: | |
| indices_in_bag.append(indices_1D[start + item_pos]) | |
| else: | |
| indices_in_bag.append(padding_idx) | |
| indices_2D[bag] = torch.tensor(indices_in_bag, device=device) | |
| return indices_2D | |
| test_cases = product(['max', 'mean', 'sum'], [False, True], [False, True], [False, True]) | |
| for mode, sparse, include_last_offset, allpad in test_cases: | |
| # Max sparse and bfloat16 are not supported | |
| if mode == 'max': | |
| if sparse or (dtype == torch.bfloat16): | |
| continue | |
| indices_1D, offsets = gen_1D_indices_offsets(include_last_offset, allpad) | |
| for padding_idx_1D in list(set(indices_1D.tolist())) + [None]: | |
| msg = ( | |
| f"mode: '{mode}', sparse: {sparse}, include_last_offset: {include_last_offset}, " | |
| f"padding_idx_1D: {padding_idx_1D}") | |
| # If 1D input does not use a padding index, we still need one for the 2D input, | |
| # so we can add one dummy word to the weights to act as the padded word | |
| padding_idx_2D = padding_idx_1D if padding_idx_1D is not None else num_words | |
| num_words_with_padding = num_words if padding_idx_1D is not None else num_words + 1 | |
| indices_2D = gen_2D_indices_from_1D( | |
| indices_1D, | |
| offsets, | |
| include_last_offset, | |
| padding_idx_2D) | |
| weights = torch.randn( | |
| num_words_with_padding, | |
| num_features, | |
| dtype=dtype, | |
| device=device, | |
| requires_grad=True) | |
| weights_check = weights.clone().detach().requires_grad_(True) | |
| bag = torch.nn.functional.embedding_bag( | |
| indices_1D, | |
| weights, | |
| offsets, | |
| padding_idx=padding_idx_1D, | |
| mode=mode, | |
| sparse=sparse, | |
| include_last_offset=include_last_offset) | |
| bag_check = torch.nn.functional.embedding_bag( | |
| indices_2D, | |
| weights_check, | |
| padding_idx=padding_idx_2D, | |
| mode=mode, | |
| sparse=sparse) | |
| self.assertEqual(bag, bag_check, msg=msg) | |
| bag.sum().backward() | |
| bag_check.sum().backward() | |
| # Sometimes, half dtype gradients mismatch by a greater amount | |
| # than other dtypes | |
| if dtype in [torch.half, torch.bfloat16]: | |
| atol = 0.01 | |
| rtol = 0.01 | |
| else: | |
| atol = None | |
| rtol = None | |
| self.assertEqual(weights.grad, weights_check.grad, msg=msg, atol=atol, rtol=rtol) | |
| # Check correctness of torch.nn.functional.embedding_bag forward and | |
| # backward functions with padding_idx, given a 2D indices input. Compare | |
| # against torch.nn.functional.embedding followed by a reduction. | |
| @onlyNativeDeviceTypes | |
| @dtypes(torch.float32, torch.float64) | |
| @dtypesIfCUDA(torch.half, torch.bfloat16) | |
| def test_embedding_bag_2D_padding_idx(self, device, dtype): | |
| # Use a Python implementation of embedding_bag with padding_idx support | |
| # to check torch.nn.functional.embedding_bag correctness | |
| def embedding_bag_check(indices, weights, mode, sparse, padding_idx): | |
| assert padding_idx is not None | |
| embedding = torch.nn.functional.embedding( | |
| indices, | |
| weights, | |
| padding_idx=padding_idx, | |
| sparse=sparse) | |
| reduction_dim = indices.dim() - 1 | |
| if mode == 'sum' or mode == 'mean': | |
| # We must avoid including elements at padding_idx in the | |
| # sum/mean, so multiply those elements by 0, and multiply | |
| # all other elements by 1 | |
| per_sample_weights = indices.ne(padding_idx).to(dtype).unsqueeze(-1) | |
| res = embedding.mul(per_sample_weights).sum(dim=reduction_dim) | |
| if mode == 'mean': | |
| weights_sum = per_sample_weights.sum(dim=reduction_dim) | |
| res = res.div(weights_sum) | |
| elif mode == 'max': | |
| # We must avoid allowing elements at padding_idx to be chosen | |
| # as the max, so set those elements to negative infinity | |
| res = embedding.masked_fill( | |
| indices.unsqueeze(-1) == padding_idx, -float('inf') | |
| ).amax(dim=reduction_dim) | |
| else: | |
| raise RuntimeError(f"mode '{mode}' is not available") | |
| # If a row is all padding, set its corresponding result row to 0. | |
| # This is needed because the above mean and max mode | |
| # implementations set these elements to nan and -inf, respectively | |
| if mode in ['mean', 'max']: | |
| res = res.masked_fill( | |
| indices.eq(padding_idx).all(dim=-1).unsqueeze(-1), | |
| 0) | |
| return res | |
| num_features = 3 | |
| num_words = 10 | |
| indices_dim1 = 10 | |
| for mode, sparse, allpad, indices_dim0 in product(['max', 'mean', 'sum'], [False, True], [False, True], [1, 10]): | |
| # Max sparse and bfloat16 are not supported | |
| if mode == 'max': | |
| if sparse or (dtype == torch.bfloat16): | |
| continue | |
| if allpad: | |
| indices = torch.empty(indices_dim0, indices_dim1, dtype=torch.long, device=device).fill_(1) | |
| else: | |
| indices = torch.randint(0, num_words, (indices_dim0, indices_dim1), device=device) | |
| if indices_dim0 > 1: | |
| # Fill one row with duplicate index so we can test with a fully | |
| # padded row | |
| duplicate_row = random.randint(0, indices_dim0 - 1) | |
| indices[duplicate_row] = indices[duplicate_row][0] | |
| for padding_idx in list(set(indices.flatten(0, -1).tolist())): | |
| weights = torch.randn(num_words, num_features, dtype=dtype, device=device, requires_grad=True) | |
| weights_check = weights.clone().detach().requires_grad_(True) | |
| msg = ( | |
| f"mode: '{mode}', sparse: {sparse}, padding_idx: {padding_idx}, " | |
| f"allpad: {allpad}, indices.size(): {indices.size()}") | |
| # Check forward with a Python implementation of padding_idx embedding_bag | |
| bag_check = embedding_bag_check( | |
| indices, | |
| weights_check, | |
| mode, | |
| sparse, | |
| padding_idx) | |
| bag = torch.nn.functional.embedding_bag( | |
| indices, | |
| weights, | |
| padding_idx=padding_idx, | |
| mode=mode, | |
| sparse=sparse) | |
| self.assertEqual(bag, bag_check, msg=msg) | |
| bag_check.sum().backward() | |
| grad_check = weights_check.grad | |
| bag.sum().backward() | |
| grad = weights.grad | |
| # Sometimes, half dtype gradients mismatch by a greater amount | |
| # than other dtypes | |
| if dtype in [torch.half, torch.bfloat16]: | |
| atol = 0.01 | |
| rtol = 0.01 | |
| else: | |
| atol = None | |
| rtol = None | |
| self.assertEqual(grad, grad_check, msg=msg, atol=atol, rtol=rtol) | |
| @onlyCUDA | |
| @dtypes(*((torch.float, torch.double, torch.bfloat16, torch.half) | |
| if TEST_WITH_ROCM else (torch.float, torch.double, torch.half))) | |
| def test_embedding_max_norm_device(self, device, dtype): | |
| embedding = nn.Embedding(22, 5, max_norm=1.0).to(device, dtype=dtype) | |
| # nn.Embedding only takes LongTensor as input | |
| input = torch.tensor([2, 8, 8, 6], device=device, dtype=torch.long) | |
| output = embedding(input) | |
| self.assertEqual(output[1], output[2]) | |
| self.assertTrue(output.data.norm(p=2, dim=1).le(1).all()) | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long))) | |
| def test_embedding_bag_empty_input(self, device, dtypes): | |
| m = 4 | |
| n = 3 | |
| x = torch.tensor([], device=device, dtype=dtypes[0]) | |
| for sparse in [True, False]: | |
| Embed = torch.nn.EmbeddingBag(m, n, sparse=sparse) | |
| Embed.to(device) | |
| output = Embed(input=x, offsets=torch.tensor([0], device=device, dtype=dtypes[1])) | |
| self.assertEqual(output, torch.zeros_like(output)) | |
| output = Embed(input=x, offsets=torch.tensor([0, 0], device=device, dtype=dtypes[1])) | |
| self.assertEqual(output, torch.zeros_like(output)) | |
| @skipCUDAIf(True, "no out-of-bounds check on CUDA for perf.") | |
| @dtypes(*itertools.product((torch.float, torch.double), (torch.int, torch.long))) | |
| @parametrize_test("padding_idx", [None, 0]) | |
| @parametrize_test("mode", ["sum", "mean", "max"]) | |
| def test_embedding_bag_out_of_bounds_idx(self, device, dtypes, padding_idx, mode): | |
| padding_idx = 0 | |
| w_dtype, idx_dtype = dtypes | |
| # negative out-of-bound | |
| idx1 = torch.tensor([[-1, 1]], device=device, dtype=idx_dtype) | |
| # positive out-of-bound | |
| idx2 = torch.tensor([[11, 8]], device=device, dtype=idx_dtype) | |
| weight = torch.randn(10, 2, device=device, dtype=w_dtype) | |
| if mode == 'sum': | |
| # Only `sum` supports per_sample_weight | |
| per_sample_weights = (None, torch.randn_like(idx1, device=device, dtype=w_dtype)) | |
| else: | |
| per_sample_weights = (None,) | |
| for p_s_weights, idx in itertools.product(per_sample_weights, (idx1, idx2)): | |
| msg = "Expected idx >= 0 && idx < num_embeddings" | |
| with self.assertRaisesRegex(RuntimeError, msg): | |
| torch.nn.functional.embedding_bag(idx, weight, | |
| per_sample_weights=p_s_weights, padding_idx=padding_idx, | |
| mode=mode) | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long))) | |
| def test_EmbeddingBag_per_sample_weights_failures(self, device, dtypes): | |
| # Failure 1: mismatched embeddings / per_sample_weights dtype | |
| es = nn.EmbeddingBag(5, 2, mode='sum').to(dtype=torch.float, device=device) | |
| input = torch.tensor([3, 1, 1, 1, 4, 0], dtype=dtypes[0], device=device) | |
| offsets = torch.tensor([0, 0, 3, 3, 6], dtype=dtypes[1], device=device) | |
| per_sample_weights = torch.randn_like(input, dtype=torch.double, device=device) | |
| if device == 'cpu': | |
| with self.assertRaisesRegex(RuntimeError, 'have the same type as'): | |
| es(input, offsets, per_sample_weights) | |
| else: | |
| with self.assertRaisesRegex(RuntimeError, 'expected scalar type'): | |
| es(input, offsets, per_sample_weights) | |
| # Failure 2.1: input/per_sample_weights have different sizes (1d input) | |
| input = torch.tensor([3, 1, 1, 1, 4, 0], dtype=dtypes[0], device=device) | |
| offsets = torch.tensor([0, 0, 3, 3, 6], dtype=dtypes[1], device=device) | |
| per_sample_weights = torch.randn(5, dtype=torch.float, device=device) | |
| with self.assertRaisesRegex(ValueError, 'same shape as the input'): | |
| es(input, offsets, per_sample_weights) | |
| # Failure 2.2: input/per_sample_weights have different sizes (2d input) | |
| input = torch.randint(5, (7, 3), dtype=dtypes[0], device=device) | |
| offsets = None | |
| per_sample_weights = torch.randn(7 * 3, dtype=torch.float, device=device) | |
| with self.assertRaisesRegex(ValueError, 'same shape as the input'): | |
| es(input, offsets, per_sample_weights) | |
| # Failure 3: Unsupported per_sample_weights and mode=('max', 'mean') | |
| for unsupported_mode in ('max', 'mean'): | |
| es = nn.EmbeddingBag(5, 2, mode=unsupported_mode).to( | |
| dtype=torch.float, device=device) | |
| input = torch.randint(5, (7, 3), dtype=dtypes[0], device=device) | |
| offsets = None | |
| per_sample_weights = torch.randn(7, 3, dtype=torch.float, device=device) | |
| with self.assertRaisesRegex(NotImplementedError, | |
| "only supported for mode='sum'"): | |
| es(input, offsets, per_sample_weights) | |
| def _embedding_bag_reference_impl(self, input, weight, offsets=None, mode='sum', | |
| per_sample_weights=None, include_last_offset=False): | |
| assert mode == 'sum' or per_sample_weights is None | |
| assert offsets is not None | |
| if per_sample_weights is None: | |
| per_sample_weights = torch.ones(input.size()).to( | |
| dtype=weight.dtype, device=weight.device | |
| ) | |
| assert input.numel() == per_sample_weights.numel() | |
| bags = [] | |
| long_input = input.to(torch.long) | |
| embeddings = weight.index_select(0, long_input) * per_sample_weights.unsqueeze(1) | |
| if include_last_offset: | |
| for index in range(len(offsets) - 1): | |
| offset = offsets[index] | |
| next_offset = offsets[index + 1] | |
| length = next_offset - offset | |
| if length == 0: | |
| bags.append( | |
| torch.tensor([0] * weight.size(1)).to( | |
| dtype=embeddings.dtype, device=embeddings.device | |
| ) | |
| ) | |
| else: | |
| if mode == 'sum': | |
| bags.append(embeddings.narrow(0, offset, length).sum(0)) | |
| elif mode == 'mean': | |
| bags.append(embeddings.narrow(0, offset, length).sum(0).div(length)) | |
| else: | |
| assert mode == 'max' | |
| bags.append(embeddings.narrow(0, offset, length).max(0)[0]) | |
| else: | |
| for index, offset in enumerate(offsets): | |
| if index + 1 < len(offsets): | |
| next_offset = offsets[index + 1] | |
| else: | |
| next_offset = len(long_input) | |
| length = next_offset - offset | |
| if length == 0: | |
| bags.append( | |
| torch.tensor([0] * weight.size(1)).to( | |
| dtype=embeddings.dtype, device=embeddings.device | |
| ) | |
| ) | |
| else: | |
| if mode == 'sum': | |
| bags.append(embeddings.narrow(0, offset, length).sum(0)) | |
| elif mode == 'mean': | |
| bags.append(embeddings.narrow(0, offset, length).sum(0).div(length)) | |
| else: | |
| assert mode == 'max' | |
| bags.append(embeddings.narrow(0, offset, length).max(0)[0]) | |
| return torch.stack(bags) | |
| @skipMeta | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long), (torch.half, torch.float, torch.double))) | |
| def test_EmbeddingBag_empty_per_sample_weights_and_offsets(self, device, dtypes): | |
| # Test empty input and per sample weight, and backward pass. There was a CUDA | |
| # invalid configuration bug (more context in #46572) | |
| def test_per_sample_weights(mode, trainable_scale): | |
| es = nn.EmbeddingBag(5, 2, mode=mode).to(dtype=dtypes[2], device=device) | |
| es.weight.data.copy_( | |
| torch.arange(1, 11, device=device).view_as(es.weight).to(dtypes[2])) | |
| input = torch.tensor([], device=device, dtype=dtypes[0]) | |
| offsets = torch.tensor([0, 0, 0, 0, 0], device=device, dtype=dtypes[1]) | |
| per_sample_weights = torch.randn_like(input, dtype=dtypes[2]) \ | |
| .requires_grad_(trainable_scale) | |
| ref_per_sample_weights = \ | |
| per_sample_weights.detach().requires_grad_(trainable_scale) | |
| reference_weights = es.weight.detach().requires_grad_() | |
| expected = self._embedding_bag_reference_impl( | |
| input, reference_weights, offsets, mode, ref_per_sample_weights) | |
| result = es(input, offsets, per_sample_weights) | |
| self.assertEqual(result, expected, atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| grad = torch.randn_like(expected) | |
| result.backward(grad) | |
| # the reference impl doesn't have grad fn for empty input; but the grad should | |
| # simply be a zero tensor | |
| ref_weights_grad = torch.zeros_like(es.weight) | |
| self.assertEqual(es.weight.grad, ref_weights_grad, | |
| atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| if trainable_scale: | |
| ref_per_sample_weights_grad = torch.empty_like(per_sample_weights) | |
| self.assertEqual(per_sample_weights.grad, ref_per_sample_weights_grad, | |
| atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| modes = ('sum',) | |
| trainable_scale = (True, False) | |
| for mode, trainable in itertools.product(modes, trainable_scale): | |
| test_per_sample_weights(mode, trainable) | |
| @skipMeta | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long), (torch.float, torch.double, torch.half))) | |
| def test_EmbeddingBag_per_sample_weights_and_offsets(self, device, dtypes): | |
| def test_per_sample_weights(mode, trainable_scale): | |
| es = nn.EmbeddingBag(5, 2, mode=mode).to(dtype=dtypes[2], device=device) | |
| es.weight.data.copy_( | |
| torch.arange(1, 11, device=device).view_as(es.weight).to(dtypes[2])) | |
| input = torch.tensor([3, 1, 1, 1, 4, 0], device=device, dtype=dtypes[0]) | |
| offsets = torch.tensor([0, 0, 3, 3, 6], device=device, dtype=dtypes[1]) | |
| per_sample_weights = torch.randn_like(input, dtype=dtypes[2]) \ | |
| .requires_grad_(trainable_scale) | |
| ref_per_sample_weights = \ | |
| per_sample_weights.detach().requires_grad_(trainable_scale) | |
| reference_weights = es.weight.detach().requires_grad_() | |
| expected = self._embedding_bag_reference_impl( | |
| input, reference_weights, offsets, mode, ref_per_sample_weights) | |
| result = es(input, offsets, per_sample_weights) | |
| self.assertEqual(result, expected, atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| grad = torch.randn_like(expected).to(dtype=dtypes[2], device=device) | |
| result.backward(grad) | |
| expected.backward(grad) | |
| self.assertEqual(es.weight.grad, reference_weights.grad, | |
| atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| if trainable_scale: | |
| self.assertEqual(per_sample_weights.grad, ref_per_sample_weights.grad, | |
| atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| modes = ('sum',) | |
| trainable_scale = (True, False) | |
| for mode, trainable in itertools.product(modes, trainable_scale): | |
| test_per_sample_weights(mode, trainable) | |
| @skipMeta | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long), (torch.float, torch.double, torch.half))) | |
| def test_EmbeddingBag_per_sample_weights_and_new_offsets(self, device, dtypes): | |
| def test_per_sample_weights_new_offsets(mode, trainable_scale, include_last_offset, has_weight=True): | |
| es = nn.EmbeddingBag(5, 2, mode=mode, include_last_offset=include_last_offset).to(dtype=dtypes[2], device=device) | |
| es.weight.data.copy_( | |
| torch.arange(1, 11, device=device).view_as(es.weight).to(dtypes[2])) | |
| input = torch.tensor([3, 1, 1, 1, 4, 0], device=device, dtype=dtypes[0]) | |
| offsets = torch.tensor([0, 0, 3, 3, 6], device=device, dtype=dtypes[1]) | |
| if include_last_offset: | |
| offsets = torch.cat((offsets, torch.tensor([input.size(0)], device=device, dtype=dtypes[1])), 0) | |
| if has_weight: | |
| per_sample_weights = torch.randn_like(input, device=device, dtype=dtypes[2]) \ | |
| .requires_grad_(trainable_scale) | |
| ref_per_sample_weights = \ | |
| per_sample_weights.detach().requires_grad_(trainable_scale) | |
| else: | |
| per_sample_weights = None | |
| ref_per_sample_weights = None | |
| reference_weights = es.weight.detach().requires_grad_() | |
| expected = self._embedding_bag_reference_impl( | |
| input, reference_weights, offsets, mode, ref_per_sample_weights, include_last_offset) | |
| result = es(input, offsets, per_sample_weights) | |
| self.assertEqual(result, expected, atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| grad = torch.randn_like(expected) | |
| result.backward(grad) | |
| expected.backward(grad) | |
| self.assertEqual(es.weight.grad, reference_weights.grad, | |
| atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| if has_weight and trainable_scale: | |
| self.assertEqual(per_sample_weights.grad, ref_per_sample_weights.grad, | |
| atol=dtype2prec_DONTUSE[dtypes[2]], rtol=0) | |
| trainable_scale = (True, False) | |
| include_last_offset = (True, False) | |
| modes = (('sum', False), ('sum', True), ('max', False), ('mean', False)) | |
| for (mode, has_weight), trainable, include_last_offset in itertools.product( | |
| modes, trainable_scale, include_last_offset | |
| ): | |
| test_per_sample_weights_new_offsets( | |
| mode, trainable, include_last_offset, has_weight | |
| ) | |
| def _test_EmbeddingBag_vs_Embedding(self, N, D, B, L, max_norm=None, | |
| mode='mean', | |
| device='cpu', | |
| wdtype=torch.float, | |
| dtype=torch.long, | |
| test_per_sample_weights=False, | |
| trainable_per_sample_weights=False, | |
| sparse=False, | |
| test_backward=True, | |
| backward_prec=None): | |
| es = nn.EmbeddingBag(N, D, mode=mode, sparse=sparse, max_norm=max_norm).to(device, wdtype) | |
| e = nn.Embedding(N, D, max_norm=max_norm).to(device, wdtype) | |
| e.weight.data.copy_(es.weight) | |
| input = torch.randint(N, (B, L), device=device, dtype=dtype) | |
| offsets = torch.arange(0, B, device=device, dtype=dtype).mul_(L) | |
| grad_output = torch.rand(B, D, device=device, dtype=wdtype) | |
| if test_per_sample_weights: | |
| # To prevent large gradients, weights should sum to 1 for each bag | |
| per_sample_weights = \ | |
| torch.randn(B, L, device=device, dtype=wdtype).softmax(dim=-1) | |
| per_sample_weights_reference = \ | |
| per_sample_weights.clone().requires_grad_(trainable_per_sample_weights) | |
| per_sample_weights.requires_grad_(trainable_per_sample_weights) | |
| output = es(input.view(-1), offsets, per_sample_weights.view(-1)) | |
| else: | |
| output = es(input.view(-1), offsets) | |
| per_sample_weights = None | |
| per_sample_weights_reference = None | |
| if mode == 'sum': | |
| if test_per_sample_weights: | |
| ref_output = (e(input) * per_sample_weights_reference.unsqueeze(-1)).sum(1) | |
| else: | |
| ref_output = e(input).sum(1) | |
| elif mode == 'mean': | |
| assert not test_per_sample_weights | |
| ref_output = e(input).mean(1) | |
| elif mode == 'max': | |
| assert not test_per_sample_weights | |
| ref_output = e(input).max(1)[0] | |
| self.assertEqual(output, ref_output, atol=dtype2prec_DONTUSE[wdtype], rtol=0) | |
| if not test_backward: | |
| return | |
| output.backward(grad_output) | |
| ref_output.backward(grad_output) | |
| es_weight_grad = es.weight.grad.data | |
| if sparse: | |
| es_weight_grad = es.weight.grad.data.to_dense() | |
| # We have more floating point error here because we are dealing with larger numbers | |
| if backward_prec is None: | |
| needed_prec = dtype2prec_DONTUSE[wdtype] * 5 | |
| else: | |
| needed_prec = backward_prec | |
| self.assertEqual(es_weight_grad, e.weight.grad, atol=needed_prec, rtol=0) | |
| if test_per_sample_weights and trainable_per_sample_weights: | |
| self.assertEqual(per_sample_weights.grad, per_sample_weights_reference.grad, | |
| atol=dtype2prec_DONTUSE[wdtype], rtol=0) | |
| @skipCUDAIf(True, "Temporarily disabled. See t54369166") | |
| @dtypesIfCUDA(*itertools.product((torch.int, torch.long), (torch.half, torch.float, torch.double))) | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.float, torch.double))) | |
| def test_EmbeddingBag_per_sample_weights_and_no_offsets(self, device, dtypes): | |
| def run_tests(mode, sparse, trainable_per_sample_weights): | |
| kwargs = dict(test_per_sample_weights=True, device=device, | |
| mode=mode, wdtype=dtypes[1], dtype=dtypes[0], sparse=sparse, | |
| trainable_per_sample_weights=trainable_per_sample_weights) | |
| # Simple case | |
| self._test_EmbeddingBag_vs_Embedding(2, 3, 5, 7, **kwargs) | |
| # B * L > 1000 | |
| self._test_EmbeddingBag_vs_Embedding(2, 5, 53, 23, **kwargs) | |
| # Large num_embedding | |
| self._test_EmbeddingBag_vs_Embedding(101, 5, 3, 7, **kwargs) | |
| # Large embedding_dim | |
| self._test_EmbeddingBag_vs_Embedding(2, 101, 3, 7, **kwargs) | |
| modes = ('sum',) | |
| sparsity = (True, False) | |
| trainable_scale = (True, False) | |
| for mode, sparse, trainable_per_sample_weights in \ | |
| itertools.product(modes, sparsity, trainable_scale): | |
| run_tests(mode, sparse, trainable_per_sample_weights) | |
| # Test CUDA Dense on half precision | |
| if device == 'cuda': | |
| modes = ('sum',) | |
| sparsity = (False,) | |
| trainable_scale = (True, False) | |
| for mode, sparse, trainable_per_sample_weights in \ | |
| itertools.product(modes, sparsity, trainable_scale): | |
| run_tests(mode, sparse, trainable_per_sample_weights) | |
| def _test_EmbeddingBag( | |
| self, | |
| device, | |
| mode, | |
| sparse, | |
| wdtype=torch.double, | |
| dtype=torch.long, | |
| odtype=torch.long, | |
| test_backward=True, | |
| ): | |
| # check a known test example | |
| es = nn.EmbeddingBag(5, 2, mode=mode, sparse=sparse).to(device, wdtype) | |
| es.weight.data.copy_(torch.arange(1, 11, device=device).view_as(es.weight).to(wdtype)) | |
| input = torch.tensor([3, 1, 1, 1, 4, 0], device=device, dtype=dtype) | |
| offsets = torch.tensor([0, 0, 3, 3, 6], device=device, dtype=odtype) | |
| grad_output = torch.tensor( | |
| [1, 2, | |
| 3, 4], device=device, dtype=wdtype).view(2, 2) | |
| grad_output_with_empty = torch.tensor( | |
| [99, 99, | |
| 1, 2, | |
| 99, 99, | |
| 3, 4, | |
| 99, 99], device=device, dtype=wdtype).view(5, 2) | |
| if mode == "sum" or mode == "mean": | |
| denominator = 1 if mode == "sum" else 3 | |
| expected_output = torch.tensor( | |
| [[13, 16], | |
| [13, 16]], device=device, dtype=wdtype) / denominator | |
| expected_output_with_empty = torch.tensor( | |
| [[0, 0], | |
| [13, 16], | |
| [0, 0], | |
| [13, 16], | |
| [0, 0]], device=device, dtype=wdtype) / denominator | |
| expected_grad_weight = torch.tensor( | |
| [[3, 4], | |
| [5, 8], | |
| [0, 0], | |
| [1, 2], | |
| [3, 4]], device=device, dtype=wdtype) / denominator | |
| elif mode == "max": | |
| expected_output = torch.tensor( | |
| [[7, 8], | |
| [9, 10]], device=device, dtype=wdtype) | |
| expected_output_with_empty = torch.tensor( | |
| [[0, 0], | |
| [7, 8], | |
| [0, 0], | |
| [9, 10], | |
| [0, 0]], device=device, dtype=wdtype) | |
| expected_grad_weight = torch.tensor( | |
| [[0, 0], | |
| [0, 0], | |
| [0, 0], | |
| [1, 2], | |
| [3, 4]], device=device, dtype=wdtype) | |
| output = es(input, offsets) | |
| output.backward(grad_output_with_empty) | |
| es_weight_grad = es.weight.grad.data | |
| if sparse: | |
| es_weight_grad = es.weight.grad.to_dense() | |
| self.assertEqual(output, expected_output_with_empty) | |
| self.assertEqual(es_weight_grad, expected_grad_weight, atol=dtype2prec_DONTUSE[wdtype], rtol=0) | |
| # check same example except as 2D (2 x 3) | |
| input = input.view(2, -1) | |
| es.zero_grad() | |
| output = es(input) | |
| output.backward(grad_output) | |
| es_weight_grad = es.weight.grad | |
| if sparse: | |
| es_weight_grad = es.weight.grad.to_dense() | |
| self.assertEqual(output, expected_output) | |
| self.assertEqual(es_weight_grad, expected_grad_weight, atol=dtype2prec_DONTUSE[wdtype], rtol=0) | |
| # test all empty bags | |
| es.zero_grad() | |
| inputs = torch.tensor([], dtype=dtype, device=device) | |
| offsets = torch.tensor([0, 0, 0, 0], dtype=odtype, device=device) | |
| es(inputs, offsets).sum().backward() | |
| dense_grad = es.weight.grad | |
| if dense_grad.is_sparse: | |
| dense_grad = dense_grad.to_dense() | |
| self.assertEqual(dense_grad, torch.zeros_like(es.weight)) | |
| # now compare EmbeddingBag vs Embedding + Sum/Mean, for constant bag length | |
| N, D, B, L = random.randint(1, 100), random.randint(1, 100), random.randint(1, 50), random.randint(1, 50) | |
| kwargs = dict(mode=mode, sparse=sparse, device=device, wdtype=wdtype, dtype=dtype, test_backward=test_backward) | |
| self._test_EmbeddingBag_vs_Embedding(N, D, B, L, **kwargs) | |
| for max_norm in (None, 3): | |
| for p in itertools.product([1, 2], repeat=4): | |
| self._test_EmbeddingBag_vs_Embedding(*p, max_norm=max_norm, **kwargs) | |
| # check that giving illegal input combos raises error | |
| es = nn.EmbeddingBag(10, 20, mode=mode, sparse=sparse) | |
| input = torch.ones(3, 4, dtype=dtype) | |
| offset = torch.arange(0, 3, dtype=odtype) | |
| self.assertRaises(ValueError, lambda: es(input, offset)) | |
| self.assertRaises(ValueError, lambda: es(input.view(-1))) | |
| offset[0] = 1 | |
| if self.device_type == "cpu": | |
| self.assertRaises(RuntimeError, lambda: es(input.view(-1), offset)) | |
| offset[0] = 0 | |
| offset[-1] = 100 | |
| self.assertRaises(RuntimeError, lambda: es(input.view(-1), offset)) | |
| @skipMeta | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long), (torch.float, torch.double, torch.half))) | |
| def test_embedding_bag_device(self, device, dtypes): | |
| with set_default_dtype(torch.double): | |
| self._test_EmbeddingBag(device, 'sum', False, wdtype=dtypes[2], dtype=dtypes[0], odtype=dtypes[1]) | |
| self._test_EmbeddingBag(device, 'mean', False, wdtype=dtypes[2], dtype=dtypes[0], odtype=dtypes[1]) | |
| self._test_EmbeddingBag(device, 'max', False, wdtype=dtypes[2], dtype=dtypes[0], odtype=dtypes[1]) | |
| test_backward = False | |
| if self.device_type == 'cuda': | |
| # see 'todo' in test_embedding_bag. | |
| test_backward = dtypes[2] is not torch.float16 | |
| elif self.device_type == 'cpu': | |
| # TODO: figure out why precision on sparse embeddings isn't the | |
| # same as for dense. | |
| test_backward = dtypes[2] is not torch.float and dtypes[2] is not torch.float16 | |
| self._test_EmbeddingBag( | |
| device, | |
| 'sum', | |
| True, | |
| wdtype=dtypes[2], | |
| dtype=dtypes[0], | |
| odtype=dtypes[1], | |
| test_backward=test_backward, | |
| ) | |
| self._test_EmbeddingBag( | |
| device, | |
| 'mean', | |
| True, | |
| wdtype=dtypes[2], | |
| dtype=dtypes[0], | |
| odtype=dtypes[1], | |
| test_backward=test_backward, | |
| ) | |
| @skipMeta | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long), (torch.float, torch.double, torch.half))) | |
| def test_embedding_bag_non_contiguous_weight(self, device, dtypes): | |
| weight_tensor = torch.randn(3, 4, dtype=dtypes[2], device=device) | |
| weight_tensor_non_contig = weight_tensor[:, :3] # This is non-contiguous strided. | |
| weight_tensor_contig = weight_tensor_non_contig.clone().contiguous() # Contig-strided. | |
| index = torch.tensor([0, 1, 2], dtype=dtypes[0], device=device) | |
| offsets = torch.tensor([0, 2], dtype=dtypes[1], device=device) | |
| for mode in ['sum', 'mean', 'max']: | |
| output_non_contig = F.embedding_bag( | |
| input=index, | |
| weight=weight_tensor_non_contig, | |
| offsets=offsets, | |
| mode=mode, | |
| ) | |
| output_contig = F.embedding_bag( | |
| input=index, | |
| weight=weight_tensor_contig, | |
| offsets=offsets, | |
| mode=mode, | |
| ) | |
| self.assertEqual(output_non_contig, output_contig) | |
| @onlyCUDA | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long))) | |
| def test_embedding_bag_bfloat16(self, device, dtypes): | |
| with set_default_dtype(torch.double): | |
| self._test_EmbeddingBag(device, 'sum', True, | |
| wdtype=torch.bfloat16, dtype=dtypes[0], | |
| odtype=dtypes[1], test_backward=True) | |
| self._test_EmbeddingBag(device, 'mean', True, | |
| wdtype=torch.bfloat16, dtype=dtypes[0], | |
| odtype=dtypes[1], test_backward=True) | |
| @onlyNativeDeviceTypes # currently fails on XLA | |
| @dtypes(*itertools.product((torch.int, torch.long), (torch.int, torch.long))) | |
| def test_embedding_bag_half(self, device, dtypes): | |
| self._test_EmbeddingBag(device, 'sum', True, wdtype=torch.float16, dtype=dtypes[0], odtype=dtypes[1], test_backward=True) | |
| instantiate_device_type_tests(TestEmbeddingNNDeviceType, globals()) | |
| instantiate_parametrized_tests(TestEmbeddingNN) | |
| if __name__ == '__main__': | |
| run_tests() |