# torch.nested

> 原文：[`pytorch.org/docs/stable/nested.html`](https://pytorch.org/docs/stable/nested.html)> 
>
> 译者：[飞龙](https://github.com/wizardforcel)
>
> 协议：[CC BY-NC-SA 4.0](http://creativecommons.org/licenses/by-nc-sa/4.0/)


## 简介

警告

PyTorch 嵌套张量的 API 处于原型阶段，将来会有变化。

嵌套张量允许用户将一组张量打包到一个单一的、高效的数据结构中。

输入张量的唯一约束是它们的维度必须匹配。

这样可以更有效地表示元数据并访问专门构建的内核。

嵌套张量的一个应用是在各种领域中表达顺序数据。传统方法是填充可变长度序列，而嵌套张量使用户可以绕过填充。在嵌套张量上调用操作的 API 与常规`torch.Tensor`没有区别，这应该允许与现有模型无缝集成，主要区别在于输入的构造。

由于这是一个原型功能，支持的操作仍然有限。但是，我们欢迎问题、功能请求和贡献。有关贡献更多信息，请参阅[此 Readme](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/nested/README.md)。

## 构造

构造很简单，只需将一组张量传递给`torch.nested.nested_tensor`构造函数。

```py
>>> 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])
 ]) 
```

数据类型、设备和是否需要梯度可以通过通常的关键字参数选择。

```py
>>> 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`可以用来保留传递给构造函数的张量的自动求导历史。有关更多信息，请参考嵌套张量构造函数和转换函数部分。

为了形成一个有效的嵌套张量，所有传递的张量需要在维度上匹配，但其他属性则不需要。

```py
>>> 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 
```

如果其中一个维度不匹配，构造函数会抛出错误。

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

请注意，传递的张量被复制到一个连续的内存块中。生成的嵌套张量分配新的内存来存储它们，并不保留引用。

目前我们只支持一级嵌套，即一个简单的、扁平的张量列表。在未来，我们可以添加对多级嵌套的支持，例如一个完全由张量列表组成的列表。请注意，对于这种扩展，保持每个条目的嵌套级别是均匀的非常重要，以便生成的嵌套张量具有明确定义的维度。如果您有这个需求，请随时提出功能请求，以便我们可以跟踪并相应地计划。

## 大小

尽管嵌套张量不支持`.size()`（或`.shape`），但如果维度 i 是规则的，它支持`.size(i)`。

```py
>>> 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 
```

如果所有维度都是规则的，嵌套张量应该在语义上与常规的`torch.Tensor`无法区分。

```py
>>> 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`允许您检索组成部分的视图。

```py
>>> 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]`不是一个副本，而是底层内存的一个切片，表示嵌套张量的第一个条目或组成部分。

## 嵌套张量构造函数和转换函数

以下函数与嵌套张量相关：

```py
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，则与列表中最左边的张量相同`torch.dtype`

+   **layout**（`torch.layout`，可选）- 返回的嵌套张量的期望布局。仅支持步进和不规则布局。默认值：如果为 None，则为步进布局。

+   **device**（`torch.device`，可选）- 返回的嵌套张量的期望设备。默认值：如果为 None，则与列表中最左边的张量相同`torch.device`

+   **requires_grad**（[*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")，可选）- 如果自动求导应记录返回的嵌套张量上的操作。默认值：`False`。

+   **pin_memory**（[*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")，可选）- 如果设置，返回的嵌套张量将分配在固定内存中。仅适用于 CPU 张量。默认值：`False`。

返回类型

*张量*

示例：

```py
>>> 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 
```

```py
torch.nested.as_nested_tensor(tensor_list, dtype=None, device=None, layout=None)
```

从`tensor_list`张量列表构造一个保留自动求导历史的嵌套张量。

注意

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

参数

**tensor_list**（*列表***[*张量**]*）- 具有相同 ndim 的张量列表

关键字参数

+   **dtype**（`torch.dtype`，可选）- 返回的嵌套张量的期望类型。默认值：如果为 None，则与列表中最左边的张量相同`torch.dtype`

+   **device**（`torch.device`，可选）- 返回的嵌套张量的期望设备。默认值：如果为 None，则与列表中最左边的张量相同`torch.device`

+   **layout**（`torch.layout`，可选）- 返回嵌套张量的期望布局。仅支持步进和不规则布局。默认值：如果为 None，则为步进布局。

返回类型

*张量*

示例：

```py
>>> 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.]) 
```

```py
torch.nested.to_padded_tensor(input, padding, output_size=None, out=None) → Tensor
```

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

警告

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

参数

**padding**（[*float*](https://docs.python.org/3/library/functions.html#float "(在 Python v3.12)")）- 尾随条目的填充值。

关键字参数

+   **output_size**（*元组**[*[*int*](https://docs.python.org/3/library/functions.html#int "(在 Python v3.12 中)")*]*）- 输出张量的大小。如果给定，必须足够大以包含所有嵌套数据；否则，将通过沿每个维度取每个嵌套子张量的最大大小来推断。

+   **out**（*张量*，可选）- 输出张量。

示例：

```py
>>> 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()` | 支持两个 3 维嵌套张量的批量矩阵乘法。 |
| `torch.nn.Linear()` | 支持 3 维嵌套输入和一个密集的 2 维权重矩阵。 |
| `torch.nn.functional.softmax()` | 支持除`dim=0`以外的所有维度的 softmax。 |
| `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 的不规则维度。 |