torch.nested¶
Introduction¶
Warning
The PyTorch API of nested tensors is in prototype stage and will change in the near future.
NestedTensor allows the user to pack a list of Tensors into a single, efficient datastructure.
The only constraint on the input Tensors is that their dimension must match.
This enables more efficient metadata representations and access to purpose built kernels.
Construction is straightforward and involves passing a list of Tensors to the constructor.
>>> a, b = torch.arange(3), torch.arange(5) + 3
>>> a
tensor([0, 1, 2])
>>> b
tensor([3, 4, 5, 6, 7])
>>> nt = torch.nested.nested_tensor([a, b])
>>> nt
nested_tensor([
tensor([0, 1, 2]),
tensor([3, 4, 5, 6, 7])
])
Data type and device can be chosen via the usual keyword arguments.
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32, device="cuda")
>>> nt
nested_tensor([
tensor([0., 1., 2.], device='cuda:0'),
tensor([3., 4., 5., 6., 7.], device='cuda:0')
])
In order to form a valid NestedTensor the passed Tensors also all need to match in dimension, but none of the other attributes need to.
>>> a = torch.randn(3, 50, 70) # image 1
>>> b = torch.randn(3, 128, 64) # image 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt.dim()
4
If one of the dimensions don’t match, the constructor throws an error.
>>> a = torch.randn(50, 128) # text 1
>>> b = torch.randn(3, 128, 64) # image 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: All Tensors given to nested_tensor must have the same dimension. Found dimension 3 for Tensor at index 1 and dimension 2 for Tensor at index 0.
Note that the passed Tensors are being copied into a contiguous piece of memory. The resulting NestedTensor allocates new memory to store them and does not keep a reference.
At this moment we only support one level of nesting, i.e. a simple, flat list of Tensors. In the future we can add support for multiple levels of nesting, such as a list that consists entirely of lists of Tensors. Note that for this extension it is important to maintain an even level of nesting across entries so that the resulting NestedTensor has a well defined dimension. If you have a need for this feature, please feel encourage to open a feature request so that we can track it and plan accordingly.
size¶
Even though a NestedTensor does not support .size() (or .shape), it supports .size(i) if dimension i is regular.
>>> a = torch.randn(50, 128) # text 1
>>> b = torch.randn(32, 128) # text 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt.size(0)
2
>>> nt.size(1)
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
RuntimeError: Given dimension 1 is irregular and does not have a size.
>>> nt.size(2)
128
If all dimensions are regular, the NestedTensor is intended to be semantically indistinguishable from a regular torch.Tensor.
>>> a = torch.randn(20, 128) # text 1
>>> nt = torch.nested.nested_tensor([a, a], dtype=torch.float32)
>>> nt.size(0)
2
>>> nt.size(1)
20
>>> nt.size(2)
128
>>> torch.stack(nt.unbind()).size()
torch.Size([2, 20, 128])
>>> torch.stack([a, a]).size()
torch.Size([2, 20, 128])
>>> torch.equal(torch.stack(nt.unbind()), torch.stack([a, a]))
True
In the future we might make it easier to detect this condition and convert seamlessly.
Please open a feature request if you have a need for this (or any other related feature for that manner).
unbind¶
unbind allows you to retrieve a view of the constituents.
>>> import torch
>>> a = torch.randn(2, 3)
>>> b = torch.randn(3, 4)
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt
nested_tensor([
tensor([[ 1.2286, -1.2343, -1.4842],
[-0.7827, 0.6745, 0.0658]]),
tensor([[-1.1247, -0.4078, -1.0633, 0.8083],
[-0.2871, -0.2980, 0.5559, 1.9885],
[ 0.4074, 2.4855, 0.0733, 0.8285]])
])
>>> nt.unbind()
(tensor([[ 1.2286, -1.2343, -1.4842],
[-0.7827, 0.6745, 0.0658]]), tensor([[-1.1247, -0.4078, -1.0633, 0.8083],
[-0.2871, -0.2980, 0.5559, 1.9885],
[ 0.4074, 2.4855, 0.0733, 0.8285]]))
>>> nt.unbind()[0] is not a
True
>>> nt.unbind()[0].mul_(3)
tensor([[ 3.6858, -3.7030, -4.4525],
[-2.3481, 2.0236, 0.1975]])
>>> nt
nested_tensor([
tensor([[ 3.6858, -3.7030, -4.4525],
[-2.3481, 2.0236, 0.1975]]),
tensor([[-1.1247, -0.4078, -1.0633, 0.8083],
[-0.2871, -0.2980, 0.5559, 1.9885],
[ 0.4074, 2.4855, 0.0733, 0.8285]])
])
Note that nt.unbind()[0] is not a, but rather a slice of the underlying memory, which represents the first entry or constituent of the NestedTensor.
Nested tensor functions¶
The following functions are related to nested tensors:
- torch.nested.nested_tensor(tensor_list, *, dtype=None, device=None, requires_grad=False, pin_memory=False)[source]¶
Constructs a nested tensor with no autograd history (also known as a “leaf tensor”, see Autograd mechanics) from
tensor_lista list of tensors.- Parameters:
tensor_list (List[Tensor]) – a list of tensors with the same ndim
- Keyword Arguments:
dtype (
torch.dtype, optional) – the desired type of returned nested tensor. Default: if None, sametorch.dtypeas leftmost tensor in the list.device (
torch.device, optional) – the desired device of returned nested tensor. Default: if None, sametorch.deviceas leftmost tensor in the listrequires_grad (bool, optional) – If autograd should record operations on the returned nested tensor. Default:
False.pin_memory (bool, optional) – If set, returned nested tensor would be allocated in the pinned memory. Works only for CPU tensors. Default:
False.
Example:
>>> a = torch.arange(3, dtype=torch.float, requires_grad=True) >>> b = torch.arange(5, dtype=torch.float, requires_grad=True) >>> nt = torch.nested.nested_tensor([a, b], requires_grad=True) >>> nt.is_leaf True
- torch.nested.as_nested_tensor(tensor_list, dtype=None, device=None)[source]¶
Constructs a nested tensor preserving autograd history from
tensor_lista list of tensors.Note
Tensors within the list are always copied by this function due to current nested tensor semantics.
- Parameters:
tensor_list (List[Tensor]) – a list of tensors with the same ndim
- Keyword Arguments:
dtype (
torch.dtype, optional) – the desired type of returned nested tensor. Default: if None, sametorch.dtypeas leftmost tensor in the list.device (
torch.device, optional) – the desired device of returned nested tensor. Default: if None, sametorch.deviceas leftmost tensor in the list
Example:
>>> a = torch.arange(3, dtype=torch.float, requires_grad=True) >>> b = torch.arange(5, dtype=torch.float, requires_grad=True) >>> nt = torch.nested.as_nested_tensor([a, b]) >>> nt.is_leaf False >>> fake_grad = torch.nested_tensor([torch.ones_like(a), torch.zeros_like(b)]) >>> nt.backward(fake_grad) >>> a.grad tensor([1., 1., 1.]) >>> b.grad tensor([0., 0., 0., 0., 0.])
- torch.nested.to_padded_tensor(input, padding, output_size=None, out=None) Tensor¶
Returns a new (non-nested) Tensor by padding the attr:input nested tensor. The leading entries will be filled with the nested data, while the trailing entries will be padded.
Warning
to_padded_tensor()always copies the underlying data, since the nested and the non-nested tensors differ in memory layout.- Parameters:
padding (float) – The padding value for the trailing entries.
- Keyword Arguments:
Example:
>>> nt = torch.nested.nested_tensor([torch.randn((2, 5)), torch.randn((3, 4))]) nested_tensor([ tensor([[ 1.6862, -1.1282, 1.1031, 0.0464, -1.3276], [-1.9967, -1.0054, 1.8972, 0.9174, -1.4995]]), tensor([[-1.8546, -0.7194, -0.2918, -0.1846], [ 0.2773, 0.8793, -0.5183, -0.6447], [ 1.8009, 1.8468, -0.9832, -1.5272]]) ]) >>> pt_infer = torch.nested.to_padded_tensor(nt, 0.0) tensor([[[ 1.6862, -1.1282, 1.1031, 0.0464, -1.3276], [-1.9967, -1.0054, 1.8972, 0.9174, -1.4995], [ 0.0000, 0.0000, 0.0000, 0.0000, 0.0000]], [[-1.8546, -0.7194, -0.2918, -0.1846, 0.0000], [ 0.2773, 0.8793, -0.5183, -0.6447, 0.0000], [ 1.8009, 1.8468, -0.9832, -1.5272, 0.0000]]]) >>> pt_large = torch.nested.to_padded_tensor(nt, 1.0, (2, 4, 6)) tensor([[[ 1.6862, -1.1282, 1.1031, 0.0464, -1.3276, 1.0000], [-1.9967, -1.0054, 1.8972, 0.9174, -1.4995, 1.0000], [ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000], [ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]], [[-1.8546, -0.7194, -0.2918, -0.1846, 1.0000, 1.0000], [ 0.2773, 0.8793, -0.5183, -0.6447, 1.0000, 1.0000], [ 1.8009, 1.8468, -0.9832, -1.5272, 1.0000, 1.0000], [ 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000]]]) >>> pt_small = torch.nested.to_padded_tensor(nt, 2.0, (2, 2, 2)) RuntimeError: Value in output_size is less than NestedTensor padded size. Truncation is not supported.
