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)
