注意
跳转至页面底部下载完整示例代码。
Nested Tensors 入门#
Nested Tensor（嵌套张量）概括了常规稠密张量的形状，允许表示不规则大小（ragged-sized）的数据。
对于常规张量，每个维度都是规则的并具有特定大小
对于嵌套张量，并非所有维度都具有规则的大小；其中一些是不规则的（锯齿状的）
嵌套张量是表示各种领域内序列数据的自然解决方案
在 NLP 中，句子可以具有不同的长度，因此一批句子形成一个嵌套张量
在计算机视觉（CV）中，图像可以具有不同的形状，因此一批图像形成一个嵌套张量
在本教程中，我们将演示嵌套张量的基本用法，并通过一个现实世界的例子说明其在处理不同长度序列数据时的实用性。特别地，它们对于构建能够高效处理不规则序列输入的 Transformer 来说非常有价值。下面，我们将展示如何使用嵌套张量实现多头注意力机制（Multi-head attention），该实现结合 torch.compile 使用时，性能优于直接在填充后的张量上进行运算。
嵌套张量目前是一个原型功能，后续可能会发生变化。
嵌套张量初始化#
在 Python 前端，可以通过张量列表创建嵌套张量。我们将 nt[i] 表示为 nestedtensor 的第 i 个张量分量。
通过将所有底层张量填充（padding）到相同的形状，nestedtensor 可以转换为常规张量。
所有张量都具有一个用于判断它们是否为嵌套张量的属性；
通常从不规则形状的张量批次中构建 nestedtensors。即维度 0 被假定为批处理维度。索引维度 0 会返回第一个底层的张量分量。
# When indexing a nestedtensor's 0th dimension, the result is a regular tensor.
需要注意的是，目前尚不支持对维度 0 进行切片。这意味着目前无法构建一个组合了底层张量分量的视图（view）。
嵌套张量运算#
由于每个操作都必须为 nestedtensors 显式实现，因此 nestedtensors 的操作覆盖范围目前比常规张量窄。目前，仅涵盖了索引、dropout、softmax、转置（transpose）、重塑（reshape）、线性层（linear）、bmm 等基本操作。然而，覆盖范围正在扩大。如果您需要特定的操作，请提交 issue 以帮助我们确定优先级。
reshape
reshape 操作用于更改张量的形状。其在常规张量上的完整语义可以在此处找到。对于常规张量，在指定新形状时，单个维度可以是 -1，在这种情况下，它会根据剩余维度和元素数量推断出来。
nestedtensors 的语义类似，只是 -1 不再进行推断。相反，它继承了旧的大小（此处 nt[0] 为 2，nt[1] 为 3）。对于锯齿状维度，-1 是唯一合法的指定大小。
转置
transpose 操作用于交换张量的两个维度。其完整语义可以在此处找到。请注意，对于 nestedtensors，维度 0 是特殊的；它被假定为批处理维度，因此不支持涉及 nestedtensor 维度 0 的转置。
其他
其他操作与常规张量具有相同的语义。将操作应用于 nestedtensor 等同于将操作应用于底层的张量分量，结果同样是一个 nestedtensor。
为什么使用嵌套张量#
当数据是序列数据时，每个样本通常具有不同的长度。例如，在一批句子中，每个句子的单词数量不同。处理变长序列的一种常见技术是手动将每个数据张量填充到相同的形状以形成批次。例如，我们有 2 个不同长度的句子和一个词汇表。为了将其表示为单个张量，我们用 0 填充到该批次中的最大长度。
这种将一批数据填充到最大长度的技术并不理想。填充后的数据在计算中是不需要的，并且通过分配比必要更大的张量浪费了内存。此外，当应用于填充数据时，并非所有操作都具有相同的语义。对于矩阵乘法，为了忽略填充项，需要用 0 填充，而对于 softmax，必须用 -inf 填充以忽略特定项。嵌套张量的主要目标是使用标准的 PyTorch 张量用户体验（UX）来促进对不规则数据的操作，从而消除对低效且复杂的填充和掩码的需求。
让我们看一个实际的例子：Transformers 中使用的多头注意力组件。我们可以实现它，使其既能处理填充后的张量，也能处理嵌套张量。
按照 Transformer 论文设置超参数
除了 dropout 概率外：为了进行正确性检查，设置为 0
让我们根据齐夫定律（Zipf’s law）生成一些真实的模拟数据。
创建嵌套张量批次输入
生成用于比较的 query、key、value 的填充形式
构建模型
检查正确性和性能
# padding-specific step: remove output projection bias from padded entries for fair comparison
# warm up compile first...
# ...now benchmark
# warm up compile first...
# ...now benchmark
# padding-specific step: remove output projection bias from padded entries for fair comparison
请注意，如果不使用 torch.compile，Python 子类嵌套张量的开销可能使其比在填充张量上进行等效计算更慢。然而，一旦启用了 torch.compile，在嵌套张量上进行操作会带来数倍的加速。随着批次中填充比例的增加，避免在填充上进行无意义计算的价值会愈发凸显。
结论#
在本教程中，我们学习了如何使用嵌套张量执行基本操作，以及如何以避免对填充部分进行计算的方式实现 Transformer 的多头注意力。更多信息，请查看 torch.nested 命名空间的文档。
另请参阅#
# %%%%%%RUNNABLE_CODE_REMOVED%%%%%%
脚本总运行时间：（0 分 0.002 秒）