def test_nn_functional_conv2d(self): tensor1 = torch.rand(3, 128, 128) tensor2 = torch.rand(3, 300, 400) inputs = [tensor1, tensor2] weight = torch.rand(3, 3, 7, 7) # no optional params tensor_res = [torch.nn.functional.conv2d( t.unsqueeze(0), weight).squeeze(0) for t in inputs] for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = [t for t in torch.nn.functional.conv2d( nt, weight).unbind()] self.assertEqual(nt_res, tensor_res) # optional params with no bias tensor_res = [torch.nn.functional.conv2d(t.unsqueeze( 0), weight, None, 2, 3, 1, 1).squeeze(0) for t in inputs] for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = [t for t in torch.nn.functional.conv2d( nt, weight, None, 2, 3, 1, 1).unbind()] self.assertEqual(nt_res, tensor_res) # optional params with bias bias = torch.rand(3) tensor_res = [torch.nn.functional.conv2d(t.unsqueeze( 0), weight, bias, (2, 2), (3, 3), (1, 1), 1).squeeze(0) for t in inputs] for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = [t for t in torch.nn.functional.conv2d( nt, weight, bias, (2, 2), (3, 3), (1, 1), 1).unbind()] self.assertEqual(nt_res, tensor_res)
def test_unbind(self): # This is the most important operation. We want to make sure # that the Tensors we use for construction can be retrieved # and used independently while still being kept track of. # In fact nestedtensor.as_nested_tensor behave just like a list. Any # list of torch.Tensors you initialize it with will be # unbound to have the same id. That is, they are indeed # the same Variable, since each torch::autograd::Variable has # assigned to it a unique PyObject* by construction. # TODO: Check that unbind returns torch.Tensors when nested_dim is 1 a = torch.tensor([1, 2]) b = torch.tensor([7, 8]) nt = nestedtensor.as_nested_tensor([a, b]) a1, b1 = nt.unbind() self.assertTrue(a is a1) self.assertTrue(b is b1) c = torch.tensor([3, 4]) d = torch.tensor([5, 6]) e = torch.tensor([6, 7]) nt1 = nestedtensor.as_nested_tensor([[c, d], [e]]) nt11, nt12 = nt1.unbind() c1, d1 = nt11.unbind() e1 = nt12.unbind()[0] self.assertTrue(c is c1) self.assertTrue(d is d1) self.assertTrue(e is e1)
def test_nn_conv2d(self): inputs = [ torch.randn(3, 500, 600), torch.randn(3, 128, 128) ] # most of optional params conv2d = torch.nn.Conv2d(3, 33, kernel_size=3, stride=(2, 1), padding=( 4, 2), padding_mode='zeros', dilation=1, groups=1, bias=True) tensor_res = [] for i in range(2): t_res = conv2d(inputs[i].unsqueeze(0).contiguous()) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = conv2d(nt) self.assertEqual(nestedtensor.nested_tensor( tensor_res, requires_grad=True), nt_res) # some of optional params conv2d = torch.nn.Conv2d(3, 33, kernel_size=3, bias=False) tensor_res = [] for i in range(2): t_res = conv2d(inputs[i].unsqueeze(0).contiguous()) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = conv2d(nt) self.assertEqual(nestedtensor.nested_tensor( tensor_res, requires_grad=True), nt_res)
def test_as_nested_tensor(self): tensors = [] num_tensors = 16 for i in range(num_tensors): tensors.append(utils.gen_float_tensor(i, (i + 1, 128, 128))) # This should NOT create references nested_tensor = nestedtensor.as_nested_tensor(tensors) for i in range(num_tensors): tensors[i].mul_(i + 2) for i in range(num_tensors): self.assertNotEqual(tensors[i], nested_tensor.unbind()[i]) # This should NOT create references nested_tensor = nestedtensor.nested_tensor(tensors) for i in range(num_tensors): tensors[i].mul_(i + 2) for i in range(num_tensors): self.assertNotEqual(tensors[i], nested_tensor.unbind()[i]) nested_tensor1 = nestedtensor.as_nested_tensor(nested_tensor) self.assertTrue(nested_tensor1 is nested_tensor) self.assertRaises( NotImplementedError, lambda: nestedtensor.as_nested_tensor(nested_tensor, dtype=torch.int64))
def test_grad_to_tensor_mask(self): def some_func(x): return torch.sum(x**2 + x**3) nt1 = nestedtensor.as_nested_tensor([ torch.tensor([1, 2, 3, 4]), torch.tensor([1, 2, 3]), torch.tensor([1, 2]) ], dtype=torch.float, requires_grad=True) nt_sum_res = some_func(nt1) nt_sum_res.backward() self.assertEqual(nt1[0].grad, torch.tensor([5., 16., 33., 56.])) self.assertEqual(nt1[1].grad, torch.tensor([5., 16., 33.])) self.assertEqual(nt1[2].grad, torch.tensor([5., 16.])) nt2 = nestedtensor.as_nested_tensor([ torch.tensor([1, 2, 3, 4]), torch.tensor([1, 2, 3]), torch.tensor([1, 2]) ], dtype=torch.float, requires_grad=True) tensor, mask = nt2.to_tensor_mask(mask_dim=2) sum_res = some_func(tensor) sum_res.backward() self.assertEqual(sum_res, nt_sum_res) self.assertEqual(nt2[0].grad, torch.tensor([5., 16., 33., 56.])) self.assertEqual(nt2[1].grad, torch.tensor([5., 16., 33.])) self.assertEqual(nt2[2].grad, torch.tensor([5., 16.]))
def test_nn_functional_interpolate(self): inputs = [ torch.randn(3, 200, 300), torch.randn(3, 300, 400) ] # no optional params tensor_res = [] for i in range(2): t_res = torch.nn.functional.interpolate( inputs[i].unsqueeze(0).contiguous(), 200) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = torch.nn.functional.interpolate(nt, 200) self.assertEqual(nestedtensor.nested_tensor(tensor_res), nt_res) # tuple/int size and optional mode for size in [(200, 200), 100]: tensor_res = [] for i in range(2): t_res = torch.nn.functional.interpolate(inputs[i].unsqueeze( 0).contiguous(), size, mode='bilinear', align_corners=True) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = torch.nn.functional.interpolate( nt, size, mode='bilinear', align_corners=True) self.assertEqual( nestedtensor.nested_tensor(tensor_res), nt_res) # special NT case - list of sizes size = ((100, 100), (200, 250), ) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = torch.nn.functional.interpolate( nt, size, mode='bilinear', align_corners=True) self.assertEqual(nt_res.nested_size(2), (100, 200)) self.assertEqual(nt_res.nested_size(3), (100, 250)) # scale_factor instead of a size for scale_factor in [(2.2, 2.2), 1.1]: tensor_res = [] for i in range(2): t_res = torch.nn.functional.interpolate( inputs[i].unsqueeze(0).contiguous(), scale_factor=scale_factor) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = torch.nn.functional.interpolate( nt, scale_factor=scale_factor) self.assertEqual( nestedtensor.nested_tensor(tensor_res), nt_res) # check errors for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: self.assertRaises(RuntimeError, lambda: torch.nn.functional.interpolate( nt, size=(100, 100), scale_factor=(1, 1)))
def test_len(self): a = nestedtensor.as_nested_tensor([torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor([5, 6]), torch.tensor([7, 8])]) self.assertEqual(len(a), 4) a = nestedtensor.as_nested_tensor([torch.tensor([1, 2]), torch.tensor([7, 8])]) self.assertEqual(len(a), 2) a = nestedtensor.as_nested_tensor([torch.tensor([1, 2])]) self.assertEqual(len(a), 1)
def test_nn_max_pool2d(self): data = [ [ torch.randn(3, 500, 600), torch.randn(3, 128, 128) ], [ torch.randn(3, 500, 600), torch.randn(3, 500, 600) ], ] # with optional params maxPool2d = torch.nn.MaxPool2d(kernel_size=( 3, 3), stride=2, padding=(1, 1), dilation=1, ceil_mode=False) for inputs in data: tensor_res = [] for i in range(2): t_res = maxPool2d(inputs[i].unsqueeze(0).contiguous()) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = maxPool2d(nt) self.assertEqual( nestedtensor.nested_tensor(tensor_res), nt_res)
def to(self, *args, **kwargs): # TODO: to is currently not supported by impls due to argparsing. new_tensors = [t.to(*args, **kwargs) for t in self.unbind()] # TODO: Make contiguous by default? Heavy operation... # NOTE: Needs grad support, which nestedtensor.nested_tensor # constructor doesn't have. return NestedTensor(nestedtensor.as_nested_tensor(new_tensors))
def test_basic_grad(self): def some_func(x): return torch.sum(x**2 + x**3) # single tensor case for comparison verification_tensor = torch.tensor([[1, 2], [3, 4]], dtype=torch.float, requires_grad=True) sum_res = some_func(verification_tensor) sum_res.backward() # as_nested_tensor constructor tensor = torch.tensor([[1, 2], [3, 4]], dtype=torch.float, requires_grad=True) nt = nestedtensor.as_nested_tensor([tensor]) nt_sum_res = some_func(nt) self.assertRaisesRegex( RuntimeError, "element 0 of tensors does not require grad and does not have a grad_fn", lambda: nt_sum_res.backward()) self.assertEqual(sum_res, nt_sum_res) self.assertIsNone(nt[0].grad) self.assertIsNotNone(verification_tensor.grad) # nested_tensor constructor tensor2 = torch.tensor([[1, 2], [3, 4]], dtype=torch.float, requires_grad=True) nt2 = nestedtensor.nested_tensor([tensor2]) #, requires_grad=True) nt_sum_res2 = some_func(nt2) # TODO: Re-enable under autograd self.assertRaises(RuntimeError, lambda: nt_sum_res2.backward())
def test_pin_memory(self): # Check if it can be applied widely nt = utils.gen_nested_tensor(1, 4, 3) nt1 = nt.pin_memory() # Make sure it's actually a copy self.assertFalse(nt.is_pinned()) self.assertTrue(nt1.is_pinned()) a1 = torch.randn(1, 2) a2 = torch.randn(2, 3) nt2 = nestedtensor.as_nested_tensor([a1, a2]) self.assertFalse(a1.is_pinned()) self.assertFalse(a2.is_pinned()) # Double check property transfers nt3 = nt2.pin_memory() self.assertFalse(nt2.is_pinned()) self.assertTrue(nt3.is_pinned()) # Check whether pinned memory is applied to constiuents # and relevant constiuents only. a3, a4 = nt3.unbind() a5, a6 = nt2.unbind() self.assertFalse(a1.is_pinned()) self.assertFalse(a2.is_pinned()) self.assertTrue(a3.is_pinned()) self.assertTrue(a4.is_pinned()) self.assertFalse(a5.is_pinned()) self.assertFalse(a6.is_pinned())
def test_nn_functional_max_pool2d(self): inputs = [torch.randn(3, 500, 600), torch.randn(3, 128, 128)] tensor_res = [] for i in range(2): t_res = torch.nn.functional.max_pool2d( inputs[i].unsqueeze(0).contiguous(), kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), dilation=(1, 1), ceil_mode=False) tensor_res.append(t_res.squeeze(0)) for nt in [ nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs) ]: nt_res = torch.nn.functional.max_pool2d(nt, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), dilation=(1, 1), ceil_mode=False) self.assertEqual(nestedtensor.nested_tensor(tensor_res), nt_res)
def _test_property(self, fn): num_nested_tensor = 3 nested_tensor_lists = [gen_nested_list(i, i, 3) for i in range(1, num_nested_tensor)] first_tensors = [get_first_tensor(ntl) for ntl in nested_tensor_lists] nested_tensors = [nestedtensor.as_nested_tensor(ntl) for ntl in nested_tensor_lists] for nested_tensor, first_tensor in zip(nested_tensors, first_tensors): self.assertEqual(fn(nested_tensor), fn(first_tensor))
def test_contiguous(self): a = nestedtensor.as_nested_tensor([ torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor([5, 6]), torch.tensor([7, 8]) ]) self.assertTrue(a.is_contiguous())
def gen_ant_unbind(): nested_tensor = nestedtensor.as_nested_tensor( [torch.rand(i, 2560) for i in RAND_INTS]) def ant(): nested_tensor.unbind() return ant
def gen_ant_unbind_2(): nested_tensor = nestedtensor.as_nested_tensor( [[torch.rand(i, 25) for i in RAND_INTS] for j in range(100)]) def ant_2(): [t.unbind() for t in nested_tensor.unbind()] return ant_2
def gen_current(): n = nestedtensor.as_nested_tensor( [torch.randn(256, 128).to(device='cuda') for _ in range(128)]) def _algorithm(): n1 = n + 1 n1.abs() return _algorithm
def _test_binary(self): a = utils.gen_float_tensor(1, (2, 3)) b = utils.gen_float_tensor(2, (2, 3)) c = utils.gen_float_tensor(3, (2, 3)) # The constructor is supposed to copy! a1 = nestedtensor.nested_tensor([a, b]) a2 = nestedtensor.nested_tensor([b, c]) a1_l = nestedtensor.as_nested_tensor([a.clone(), b.clone()]) a2_l = nestedtensor.as_nested_tensor([b.clone(), c.clone()]) a3 = nestedtensor.nested_tensor( [getattr(torch, func)(a, b), getattr(torch, func)(b, c)]) a3_l = nestedtensor.as_nested_tensor(a3) self.assertEqual(a3_l, getattr(torch, func)(a1_l, a2_l)) self.assertEqual(a3_l, getattr(torch, func)(a1, a2)) self.assertEqual(a3, getattr(a1, func)(a2)) self.assertEqual(a3, getattr(a1, func + "_")(a2)) self.assertEqual(a3, a1)
def gen_algorithm_nested_jit_mv(keys, sub_clusters): # For-loops over vectors and matrices new_sub_clusters = [] for sub_cluster in sub_clusters: new_sub_cluster = [] for cluster in sub_cluster: new_sub_cluster.append(torch.stack(cluster)) new_sub_clusters.append(new_sub_cluster) nested_sub_clusters = nestedtensor.as_nested_tensor(new_sub_clusters) nested_keys = nestedtensor.as_nested_tensor(keys) @nestedtensor._C.jit_tensorwise() @torch.jit.script def my_fun(x, y): return torch.mv(x, y) def _nested_jit_mv(): return my_fun(nested_sub_clusters, nested_keys) return _nested_jit_mv
def test_nn_functional_cross_entropy(self): inputs = [ torch.randn(3, 300, 300), torch.randn(3, 400, 400) ] targets = [ torch.randint(1, (300, 300), dtype=torch.int64), torch.randint(1, (400, 400), dtype=torch.int64) ] tensor_res = [] for i in range(2): t_res = torch.nn.functional.cross_entropy( inputs[i].unsqueeze(0).contiguous(), targets[i].unsqueeze(0)) tensor_res.append(t_res.squeeze(0)) for input_nt, target_nt in [(nestedtensor.nested_tensor(inputs), nestedtensor.nested_tensor(targets)), (nestedtensor.as_nested_tensor(inputs), nestedtensor.as_nested_tensor(targets))]: nt_res = torch.nn.functional.cross_entropy(input_nt, target_nt) self.assertEqual(nestedtensor.nested_tensor(tensor_res), nt_res)
def test_constructor(self): """ This tests whether nestedtensor.as_nested_tensor stores Variables that share storage with the input Variables used for construction. """ tensors = [] num_tensors = 16 for i in range(num_tensors): tensors.append(gen_float_tensor(i, (i + 1, 128, 128))) nested_tensor = nestedtensor.as_nested_tensor(tensors) for i in range(num_tensors): tensors[i].mul_(i + 2) for i in range(num_tensors): self.assertEqual(tensors[i], nested_tensor.unbind()[i]) self.assertEqual(tensors[i].storage().data_ptr(), nested_tensor.unbind()[i].storage().data_ptr())
def test_nn_functional_relu(self): inputs = [torch.randn(3, 500, 600), torch.randn(3, 128, 128)] tensor_res = [] for i in range(2): t_res = torch.nn.functional.relu( inputs[i].unsqueeze(0).contiguous()) tensor_res.append(t_res.squeeze(0)) for nt in [ nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs) ]: nt_res = torch.nn.functional.relu(nt) self.assertEqual(nestedtensor.nested_tensor(tensor_res), nt_res)
def test_default_constructor(self): # self.assertRaises(TypeError, lambda: torch.nested_tensor()) # nested_dim is 1 and dim is 1 too. default_nested_tensor = nestedtensor.as_nested_tensor([]) default_tensor = torch.tensor([]) self.assertEqual(default_nested_tensor.nested_dim(), 1) # TODO: Return torch.NestedSize instead of list self.assertEqual(nested_size_to_list(default_nested_tensor.nested_size()), []) self.assertEqual(default_nested_tensor.dim(), default_tensor.dim()) self.assertEqual(default_nested_tensor.layout, default_tensor.layout) self.assertEqual(default_nested_tensor.device, default_tensor.device) self.assertEqual(default_nested_tensor.dtype, default_tensor.dtype) self.assertEqual(default_nested_tensor.requires_grad, default_tensor.requires_grad) self.assertEqual(default_nested_tensor.is_pinned(), default_tensor.is_pinned())
def test_nn_dropout(self): inputs = [torch.randn(3, 128, 128), torch.randn(3, 300, 400)] dropout = torch.nn.Dropout(p=0.2) tensor_res = [] for i in range(2): t_res = dropout(inputs[i].unsqueeze(0).contiguous()) tensor_res.append(t_res.squeeze(0)) for nt in [ nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs) ]: nt_res = dropout(nt) self.assertEqual( nestedtensor.nested_tensor(tensor_res).size(), nt_res.size())
def test_nn_functional_dropout(self): inputs = [ torch.randn(3, 128, 128), torch.randn(3, 300, 400) ] tensor_res = [] for i in range(2): t_res = torch.nn.functional.dropout( inputs[i].unsqueeze(0).contiguous()) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = torch.nn.functional.dropout(nt) torch.nn.functional.dropout(nt, inplace=True) self.assertEqual(nestedtensor.nested_tensor( tensor_res).size(), nt_res.size())
def test_contiguous(self): for _ in range(1, 10): # data = gen_nested_list(1, 2, 3, size_low=1, size_high=3) data = [[torch.rand(1, 2), torch.rand(3, 4)], [torch.rand(5, 6)]] nt = nestedtensor.nested_tensor(data) self.assertTrue(nt.is_contiguous()) # buf = nt.flatten() self.assertEqual(nt, nt) a = nt + nt nt.cos_() nt.cos() a = nestedtensor.as_nested_tensor([torch.tensor([1, 2]), torch.tensor([3, 4]), torch.tensor([5, 6]), torch.tensor([7, 8])]) self.assertTrue(a.is_contiguous())
def test_nn_batch_norm(self): inputs = [ torch.tensor([[[-0.5000]], [[0.5000]]]), torch.tensor([[[-1.0000, 1.0000], [-0.2500, -0.5000]], [[0.2500, 0.5000], [1.5000, -1.5000]]]) ] batch_norm = torch.nn.BatchNorm2d(2, 1e-05, 0.1) batch_norm = batch_norm.eval() tensor_res = [] for i in range(2): t_res = batch_norm(inputs[i].unsqueeze(0).contiguous()) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = batch_norm(nt) self.assertEqual(nestedtensor.nested_tensor(tensor_res, requires_grad=True), nt_res)
def test_nn_functional_batch_norm(self): inputs = [ torch.tensor([[[-0.5000]], [[0.5000]]]), torch.tensor([[[-1.0000, 1.0000], [-0.2500, -0.5000]], [[0.2500, 0.5000], [1.5000, -1.5000]]]) ] tensor_res = [] running_mean = torch.rand(2) running_var = torch.rand(2) for i in range(2): t_res = torch.nn.functional.batch_norm( inputs[i].unsqueeze(0).contiguous(), running_mean, running_var) tensor_res.append(t_res.squeeze(0)) for nt in [nestedtensor.nested_tensor(inputs), nestedtensor.as_nested_tensor(inputs)]: nt_res = torch.nn.functional.batch_norm( nt, running_mean, running_var) self.assertEqual(nestedtensor.nested_tensor(tensor_res), nt_res)
def test_functional_relu_(self): orig_t1 = torch.tensor([-2, -1, 0, 1, 2]) expected_t = torch.tensor([0, 0, 0, 1, 2]) expected_nt = nestedtensor.nested_tensor([expected_t]) t_clone = orig_t1.clone() torch.nn.functional.relu_(t_clone) self.assertEqual(t_clone, expected_t) t_clone = orig_t1.clone() nt1 = nestedtensor.nested_tensor([t_clone]) torch.nn.functional.relu_(nt1) self.assertEqual(nt1, expected_nt) self.assertEqual(t_clone, orig_t1) t_clone = orig_t1.clone() nt1 = nestedtensor.as_nested_tensor([t_clone]) torch.nn.functional.relu_(nt1) self.assertEqual(nt1, expected_nt) self.assertNotEqual(t_clone, expected_t)
def to(self, *args, **kwargs): raise NotImplementedError( "NestedTensor.to is currently not implemented.") return nestedtensor.as_nested_tensor(new_tensors)