torch.nested

介绍

警告

PyTorch 中嵌套张量的 API 目前处于原型阶段，很快就会有所变化。

NestedTensor 允许用户将张量列表打包到一个高效的数据结构中。

输入张量的唯一要求是它们的维度必须一致。

这使元数据表示更高效，并能访问专门为该目的设计的内核。

NestedTensors的一个应用是表达不同领域的序列数据。虽然传统方法是对变长序列进行填充，但NestedTensor允许用户绕过这一步骤。在嵌套张量上调用操作的API与常规torch.Tensor相同，这使得它可以无缝地与现有模型集成，主要区别在于输入构建。

由于这是一个原型功能，支持的操作仍然有限。然而，我们欢迎问题报告、功能请求和贡献。关于如何贡献的更多信息可以在此 Readme中找到。

建造

构建过程非常简单，只需将张量列表传递给 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 允许你获取构成部分的视图。

>>> 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（一个张量列表）构建一个没有自动求梯度历史记录的嵌套张量（也称为“叶张量”，详见自动微分机制）。

参数

tensor_list (List[array_like]) – 张量的列表，或者是可以传递给 torch.tensor 的任何内容。
维度。 （列表中的每个元素都具有相同的） –

关键字参数

dtype (torch.dtype, 可选) – 返回的嵌套张量的数据类型。默认值：如果为 None，则与列表中最左边张量的数据类型相同。
layout (torch.layout, 可选) – 指定返回的嵌套张量的布局。支持的布局包括 strided 和 jagged，默认为 strided 布局。
device (torch.device, 可选) – 指定返回的嵌套张量所在的设备。默认情况下，如果未指定，则使用列表中最左边张量的torch.device
requires_grad (bool, 可选) – 是否应记录返回的嵌套张量上的自动求梯度操作。默认值： False。
pin_memory (bool, 可选) – 如果启用，返回的嵌套张量将分配在固定内存中。仅适用于 CPU 张量。默认值： False。

返回类型

Tensor

示例：

>>> 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)[源代码]

从张量或张量列表/元组构建一个保留自动求导历史的嵌套张量。

如果传递了嵌套张量，并且其设备、数据类型或布局与预期不符，则会进行复制并返回。需要注意的是，目前该函数不支持更改布局。

如果传递的是一个非嵌套张量，它会被视为一个大小一致的组件批处理。如果传递的设备或数据类型与输入不同，或者输入是不连续的，那么会产生一个新的副本。否则，会直接使用输入的存储。

如果提供了一个张量列表，在构建嵌套张量的过程中，列表中的张量总会被复制。

参数

ts (Tensor 或 List[Tensor] 或 Tuple[Tensor]) – 一个要作为嵌套张量处理的张量，或者具有相同 ndim 的张量列表或元组。

关键字参数

dtype (torch.dtype, 可选) – 返回的嵌套张量的数据类型。默认值：如果为 None，则与列表中最左边张量的数据类型相同。
device (torch.device, 可选) – 指定返回的嵌套张量所在的设备。默认情况下，如果未指定，则使用列表中最左边张量的torch.device
layout (torch.layout, 可选) – 指定返回的嵌套张量的布局。支持的布局包括 strided 和 jagged，默认为 strided 布局。

返回类型

Tensor

示例：

>>> 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嵌套张量，返回一个新的（非嵌套）Tensor。前导条目将使用嵌套数据填充，而尾随条目则进行填充。

警告

to_padded_tensor() 总是会复制底层数据，因为嵌套和非嵌套张量的内存布局不同。

参数

padding (浮点数) – 尾部条目的填充值。

关键字参数

output_size (Tuple[int]) – 输出张量的大小。如果提供，则必须足够大以包含所有嵌套数据；否则，将通过沿每个维度取每个嵌套子张量的最大尺寸来推断。
out (Tensor, 可选) – 指定输出张量。

示例：

>>> 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.

支持的操作

在本节中，我们总结了目前支持的嵌套张量操作及其约束条件。

PyTorch操作	约束条件
`torch.matmul()`	支持两个（>= 3维）嵌套张量之间的矩阵乘法，要求这两个张量的最后一两维为矩阵维度，并且它们前面的批处理维度大小相等（当前不支持批处理维度的广播功能）。
`torch.bmm()`	支持两个三维嵌套张量的批量矩阵乘法。
`torch.nn.Linear()`	支持三维嵌套输入和密集的二维权重矩阵。
`torch.nn.functional.softmax()`	支持在所有维度上进行softmax操作，除了dim=0维度。
`torch.nn.Dropout()`	行为与普通张量一样。
`torch.Tensor.masked_fill()`	行为与普通张量一样。
`torch.relu()`	行为与普通张量一样。
`torch.gelu()`	行为与普通张量一样。
`torch.silu()`	行为与普通张量一样。
`torch.abs()`	行为与普通张量一样。
`torch.sgn()`	行为与普通张量一样。
`torch.logical_not()`	行为与普通张量一样。
`torch.neg()`	行为与普通张量一样。
`torch.sub()`	支持两个嵌套张量之间的元素级减法运算。
`torch.add()`	支持两个嵌套张量的元素级相加，也支持标量与嵌套张量之间的相加。
`torch.mul()`	支持两个嵌套张量的逐元素乘法，也支持将嵌套张量与标量相乘。
`torch.select()`	支持在所有维度上进行选择。
`torch.clone()`	行为与普通张量一样。
`torch.detach()`	行为与普通张量一样。
`torch.unbind()`	仅支持沿`dim=0`方向的解绑。
`torch.reshape()`	支持保持`dim=0`尺寸不变的重塑（即嵌套张量的数量不能改变）。与常规张量不同，在这里`-1`表示继承现有大小。特别是，不规则维度的有效大小只能是`-1`。由于目前尚未实现大小推断功能，因此对于新维度的大小不能设置为`-1`。
`torch.Tensor.reshape_as()`	与`reshape`相似的约束。
`torch.transpose()`	支持除了`dim=0`以外所有维度的转置。
`torch.Tensor.view()`	新形状的规则与`reshape`相似。
`torch.empty_like()`	行为类似于常规张量；返回一个新的空嵌套张量（其值尚未初始化），该张量与输入的嵌套结构相匹配。
`torch.randn_like()`	行为类似于常规张量；返回一个新的嵌套张量，其中的值会根据与输入结构相匹配的标准正态分布进行随机初始化。
`torch.zeros_like()`	行为类似于常规张量；返回一个全新的嵌套张量，其中所有值都为零，并且其结构与输入张量的嵌套结构相同。
`torch.nn.LayerNorm()`	`normalized_shape` 参数不能扩展到 NestedTensor 的不规则维度。