torch.nested¶
简介¶
警告
嵌套张量的 PyTorch API 处于原型阶段,将在不久的将来发生变化。
NestedTensor 允许用户将张量列表打包到一个高效的数据结构中。
输入张量的唯一约束是它们的维度必须匹配。
这使得更有效的元数据表示和对专用内核的访问成为可能。
NestedTensor 的一个应用是在各种领域表达顺序数据。虽然传统方法是填充可变长度序列,但 NestedTensor 允许用户绕过填充。在嵌套张量上调用操作的 API 与普通 torch.Tensor
的 API 没有区别,这应该允许与现有模型无缝集成,主要区别在于 输入的构造。
由于这是一个原型特性,支持的操作 仍然有限。但是,我们欢迎问题、功能请求和贡献。有关贡献的更多信息可以在 此自述文件 中找到。
构造¶
构造很简单,涉及将张量列表传递给 torch.nested.nested_tensor
构造函数。
>>> 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])
])
数据类型、设备和是否需要梯度可以通过通常的关键字参数进行选择。
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32, device="cuda", requires_grad=True)
>>> nt
nested_tensor([
tensor([0., 1., 2.], device='cuda:0', requires_grad=True),
tensor([3., 4., 5., 6., 7.], device='cuda:0', requires_grad=True)
], device='cuda:0', requires_grad=True)
类似于 torch.as_tensor
,torch.nested.as_nested_tensor
可用于保留传递给构造函数的张量的自动梯度历史。有关更多信息,请参阅关于 嵌套张量构造函数和转换函数 的部分。
为了形成一个有效的 NestedTensor,所有传递的张量都需要在维度上匹配,但其他任何属性都不需要匹配。
>>> 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
如果其中一个维度不匹配,构造函数将抛出错误。
>>> 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.
请注意,传递的张量将被复制到一个连续的内存块中。生成的 NestedTensor 将分配新的内存来存储它们,并且不会保留引用。
目前我们只支持一层嵌套,即张量的简单扁平列表。将来我们可以添加对多层嵌套的支持,例如一个完全由张量列表组成的列表。请注意,对于此扩展,重要的是在所有条目中保持均匀的嵌套级别,以便生成的 NestedTensor 具有明确定义的维度。如果您需要此功能,请随时提交功能请求,以便我们对其进行跟踪并相应地进行规划。
大小¶
即使 NestedTensor 不支持 .size()
(或 .shape
),如果维度 i 是规则的,它也支持 .size(i)
。
>>> 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
如果所有维度都是规则的,则 NestedTensor 旨在在语义上与普通 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
将来我们可能会使其更容易检测此条件并无缝转换。
如果您需要此功能(或任何其他相关的功能),请提交功能请求。
unbind¶
unbind
允许您检索成分的视图。
>>> 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]])
])
请注意,nt.unbind()[0]
不是副本,而是底层内存的切片,它表示 NestedTensor 的第一个条目或成分。
嵌套张量构造函数和转换函数¶
以下函数与嵌套张量相关
- torch.nested.nested_tensor(tensor_list, *, dtype=None, layout=None, device=None, requires_grad=False, pin_memory=False)[源代码]¶
从
tensor_list
(一个张量列表)构造一个没有自动梯度历史的嵌套张量(也称为“叶子张量”,参见 Autograd 机制)。- 参数
tensor_list (List[array_like]) – 张量列表,或任何可以传递给 torch.tensor 的内容,
维度。 (其中列表的每个元素都具有相同的) –
- 关键字参数
dtype (
torch.dtype
,可选) – 返回的嵌套张量的期望类型。默认值:如果为 None,则与列表中左侧最远端张量的torch.dtype
相同。layout (
torch.layout
,可选) – 返回的嵌套张量的期望布局。仅支持连续和交错布局。默认值:如果为 None,则为连续布局。device (
torch.device
,可选) – 返回的嵌套张量的期望设备。默认值:如果为 None,则与列表中左侧最远端张量的torch.device
相同。requires_grad (bool, 可选) – 是否应记录返回的嵌套张量上的操作的自动梯度。默认值:
False
。pin_memory (bool, 可选) – 如果设置,则返回的嵌套张量将在固定内存中分配。仅适用于 CPU 张量。默认值:
False
。
- 返回类型
示例
>>> 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(ts, dtype=None, device=None, layout=None)[source]¶
从张量或张量列表/元组构造一个保留自动梯度历史的嵌套张量。
如果传递了一个嵌套张量,则将直接返回它,除非设备/数据类型/布局不同。请注意,转换设备/数据类型会导致复制,而此函数目前不支持转换布局。
如果传递了一个非嵌套张量,则将其视为一批大小一致的组成部分。如果传递的设备/数据类型与输入不同,或者输入是非连续的,则会产生复制。否则,将直接使用输入的存储。
如果提供了张量列表,则在构造嵌套张量期间始终复制列表中的张量。
- 参数
ts (Tensor 或 List[Tensor] 或 Tuple[Tensor]) – 将作为嵌套张量处理的张量,或者具有相同 ndim 的张量列表/元组
- 关键字参数
dtype (
torch.dtype
,可选) – 返回的嵌套张量的期望类型。默认值:如果为 None,则与列表中左侧最远端张量的torch.dtype
相同。device (
torch.device
,可选) – 返回的嵌套张量的期望设备。默认值:如果为 None,则与列表中左侧最远端张量的torch.device
相同。layout (
torch.layout
,可选) – 返回的嵌套张量的期望布局。仅支持连续和交错布局。默认值:如果为 None,则为连续布局。
- 返回类型
示例
>>> 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.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.]) >>> c = torch.randn(3, 5, requires_grad=True) >>> nt2 = torch.nested.as_nested_tensor(c)
- torch.nested.to_padded_tensor(input, padding, output_size=None, out=None) Tensor ¶
通过填充
input
嵌套张量返回一个新的(非嵌套)张量。前导条目将填充嵌套数据,而尾随条目将填充。警告
to_padded_tensor()
始终复制底层数据,因为嵌套张量和非嵌套张量的内存布局不同。- 参数
padding (float) – 尾随条目的填充值。
- 关键字参数
示例
>>> 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.
支持的操作¶
在本节中,我们总结了当前在 NestedTensor 上支持的操作及其任何约束。
PyTorch 操作 |
约束 |
---|---|
支持两个(>= 3D)嵌套张量之间的矩阵乘法,其中最后两个维度是矩阵维度,前导(批处理)维度具有相同的大小(即尚未支持批处理维度的广播)。 |
|
支持两个 3D 嵌套张量的批处理矩阵乘法。 |
|
支持 3D 嵌套输入和密集的 2D 权重矩阵。 |
|
支持除 dim=0 之外的所有维度上的 softmax。 |
|
行为与普通张量相同。 |
|
行为与普通张量相同。 |
|
|
行为与普通张量相同。 |
|
行为与普通张量相同。 |
|
行为与普通张量相同。 |
行为与普通张量相同。 |
|
行为与普通张量相同。 |
|
行为与普通张量相同。 |
|
行为与普通张量相同。 |
|
支持两个嵌套张量的逐元素减法。 |
|
支持两个嵌套张量的逐元素加法。支持将标量添加到嵌套张量。 |
|
支持两个嵌套张量的逐元素乘法。支持将嵌套张量乘以标量。 |
|
支持沿所有维度选择。 |
|
行为与普通张量相同。 |
|
|
行为与普通张量相同。 |
仅支持沿 |
|
支持保留 |
|
与 |
|
支持除 |
|
新形状的规则类似于 |
|
行为类似于普通张量;返回一个新的空嵌套张量(即具有未初始化的值),与输入的嵌套结构匹配。 |
|
行为类似于普通张量;返回一个新的嵌套张量,其值根据标准正态分布随机初始化,与输入的嵌套结构匹配。 |
|
行为类似于普通张量;返回一个新的嵌套张量,所有值为零,与输入的嵌套结构匹配。 |
|
|