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可以用于保存传递给构造函数的张量的autograd历史。更多详细信息，请参阅嵌套张量构造函数和转换函数部分。

为了形成一个有效的 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 张量列表构造一个没有自动求导历史的嵌套张量（也称为“叶张量”，请参阅 自动求导机制)。

Parameters

• tensor_list (List[array_like]) – 一个张量列表，或者任何可以传递给 torch.tensor 的内容,

• 维度。 (其中每个元素为列表中的元素具有相同的) –

Keyword Arguments

• dtype (torch.dtype, 可选) – 返回的嵌套张量所需的类型。 默认值: 如果为 None，则与列表中最左侧张量的 torch.dtype 相同。

• 布局 (torch.layout, 可选) – 返回嵌套张量所需的布局。 仅支持步幅和锯齿状布局。默认值: 如果为 None，则使用步幅布局。

• 设备 (torch.device, 可选) – 返回嵌套张量所需的设备. 默认值: 如果为 None，则与列表中最左侧张量的 torch.device 相同

• requires_grad (bool, optional) – 如果自动求导应该记录对返回嵌套张量的操作。默认值: False。

• pin_memory (bool, optional) – 如果设置，返回的嵌套张量将分配在固定内存中。仅适用于CPU张量。默认值: False。

Return type

张量

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, layout=None)

从tensor_list一个张量列表构建保留自动梯度历史的嵌套张量。

注意

列表中的张量由于当前嵌套张量语义的原因，总是会被此函数复制。

Parameters

tensor_list (List[Tensor]) – 一个包含相同ndim的张量列表

Keyword Arguments

• dtype (torch.dtype, 可选) – 返回的嵌套张量所需的类型。 默认值: 如果为 None，则与列表中最左侧张量的 torch.dtype 相同。

• 设备 (torch.device, 可选) – 返回嵌套张量所需的设备. 默认值: 如果为 None，则与列表中最左侧张量的 torch.device 相同

• 布局 (torch.layout, 可选) – 返回嵌套张量所需的布局。 仅支持步幅和锯齿状布局。默认值: 如果为 None，则使用步幅布局。

Return type

张量

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

通过填充input嵌套张量，返回一个新的（非嵌套）张量。 前导条目将用嵌套数据填充， 而尾随条目将被填充。

警告

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

Parameters

填充 (浮点数) – 后续条目的填充值。

Keyword Arguments

• output_size (Tuple[int]) – 输出张量的大小。 如果提供，它必须足够大以包含所有嵌套数据； 否则，将通过取每个维度下每个嵌套子张量的最大尺寸进行推断。

• 输出 (张量, 可选) – 输出张量。

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.

支持的操作

在本节中，我们总结了目前支持的 NestedTensor 操作及其相关的限制条件。

PyTorch 操作	约束条件
torch.matmul()	支持两个（>= 3d）嵌套张量之间的矩阵乘法，其中最后两个维度是矩阵维度，前面的（批处理）维度大小相同（即目前不支持对批处理维度进行广播）。
torch.bmm()	支持两个三维嵌套张量的批量矩阵乘法。
torch.nn.Linear()	支持3维嵌套输入和密集的2维权重矩阵。
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 的不规则维度中。