张量操作的分类法

张量操作是深度学习框架的构建模块，但并非所有操作都生而平等。在这篇文章中，我将为您提供一个简单的框架来理解它们，并展示为什么重塑（reshape）几乎是免费的，而矩阵乘法则会消耗大量的GPU周期。

视图操作与数据操作

张量操作主要分为两类：视图操作（view operations），它们只改变张量的索引方式，而不移动数据；以及数据操作（data operations），它们改变底层的值。视图操作通常非常快，因为它们只改变索引指令，而数据操作则进行实际的计算工作。

视图操作

视图操作接受一个张量并改变其索引方式。这些操作在数学上不是很有趣，但可以使张量程序在概念上更简单或计算上更高效。

置换 (Permute)

视图操作最简单的例子是置换（permute）。

例如：

a b c -> b a c

如果你尝试访问输出张量中 [i, j, k] 的元素，它等同于访问原始张量中 [j, i, k] 的元素。

举个具体例子

表示一批RGB图像的张量形状可能是 [batch, height, width, channels]。将其置换为 [batch, channels, height, width] 重新排列了维度，而无需复制任何像素数据。

重塑 (Reshape)

在我们开始重塑之前，首先回顾几个关键概念。

展平 (Raveling) 和展开 (Unraveling)

想象一下将一个二维电子表格展平为单列。你会逐行地将每个单元格按顺序放置。这基本上就是任何多维张量存储在计算机线性内存中时发生的情况。它需要被“展平”成一维表示。

与我上一篇文章类似，我使用基于1的张量索引，因为我喜欢它从列表末尾或开头计数的对称性。然而，出于某些原因，索引向量和线性索引的实际值是基于0的。

对于具有 `shape` 的多维 `index`，这种映射的规范方式是：

线性索引 = index[-1] + index[-2]*shape[-1] + index[-3]*shape[-1]*shape[-2] + ....

例如，我们可以将灰度图像视为形状为 [h, w] 的二维张量。你可以想象这个线性索引按行向下遍历图像。

其中 i 表示高度（h），j 表示宽度（w）

线性索引 = j + i*w = i*w + j

所以每个 i 跳过一行，每个 j 在行内选择。

这被称为展平（raveling）。

如果我们想从 linear_index 中取回 index，假设我们有原始 shape，这会稍微棘手一些。如果你好奇，值得先自己尝试一下。

让我们用灰度图像的例子来澄清，其中

linear_index = i*w + j

如果我们要从中获取 j，我们应该记住如何使用模数（%）来获取除以 w 的余数。因为 j 在 [0, w) 之间

% 是取模运算符（返回除法后的余数）

j = linear_index % w

现在我们只需要 i，它等于

i = linear_index - j = linear_index - linear_index % w = linear_index // w

其中 // 是整除，即除法后向下取整到最近的整数。

对于更高维度的索引，模式如下：

index[-1] = linear_index % shape[-1]

这确保了我们 `linear_index` 中乘以 `shape[-1]` 的任何项都变为 `0`。

所以

• = linear_index % shape[-1]
• = (index[-1] + index[-2]*shape[-1] + index[-3]*shape[-1]*shape[-2] + ...) % shape[-1]
• = index[-1] % shape[-1] + index[-2]*shape[-1] % shape[-1] + index[-3]*shape[-1]*shape[-2] % shape[-1] + ...
• = index[-1] % shape[-1] + 0 + 0 + ...
• = index[-1]

然后对于下一个，我们使用前面提到的整除 // 减去之前找到的索引

• 线性索引 // shape[-1]
• = (index[-1] + index[-2]*shape[-1] + index[-3]*shape[-1]*shape[-2] + ...) // shape[-1]
• = index[-1] // shape[-1] + index[-2]*shape[-1] // shape[-1] + index[-3]*shape[-1]*shape[-2] // shape[-1] + ...
• = 0 + index[-2] + index[-3]*shape[-2] + ...
• = index[-2] + index[-3]*shape[-2] + ...

但由于我们有更多维度，我们需要再次对下一个维度取模

• 索引[-2] = (线性索引 // 形状[-1]) % 形状[-2]
• 索引[-3] = (线性索引 // 形状[-1]*形状[-2]) % 形状[-3]
• 索引[-4] = (线性索引 // 形状[-1]*形状[-2]*形状[-3]) % 形状[-4]

回到我们的灰度图像示例

如果对于一个 4x5 的图像，我们有 linear_index = 15，那么我们得到 j = 15 % 5 = 0 和 i = 15 // 5 = 3，所以我们位于 [3, 0] 的位置，即第四行的第一个像素。

这种简洁的数学运算只有在基于0的索引下才可能实现。

这被称为展开（unraveling）。

对我来说，这个命名感觉有点反直觉，因为将张量线性化感觉像是从织物中解开一根线，但无所谓了。

合并与分割 (Merge and Split)

现在，既然已经讲完了，还有两种基本的重塑操作：合并（Merge）和分割（Split）。

合并 (Merge)

合并是两者中较简单的一个，它接受一个子张量，并将多维子张量的索引转换为线性索引。

例如：

b h w c -> b (h w) c

对输出张量 [i, j, k] 进行索引，等同于对原始张量中的展开索引 [i, j // w, j % w, k] 进行索引。

合并除了相邻维度之外不需要额外信息。

分割 (Split)

分割操作更具挑战性，因为你必须知道除了一个维度之外的所有维度。

b (h w) c -> b h w c

在这种情况下，你必须知道 h 的大小或 w 的大小。

大多数张量框架在分割时都要求你明确指定目标维度，因为目标形状无法从输入中完全推断出来。

对输出张量 [i, j, k, l] 进行索引，等同于对原始张量中的展平索引 [i, j*w + k, l] 进行索引。

请注意，通过以下操作添加单例维度是多么容易：

(1 1 1 1 1 1 1 1 1 b) h c -> 1 1 1 1 1 1 1 1 1 b h c

可以看到线性索引不受影响

• 线性索引 = 索引[-1] + 索引[-2]*形状[-1] + 索引[-3]*形状[-1]*形状[-2] + 索引[-3]*形状[-1]*形状[-2]*形状[-3] + ...
• 线性索引 = k + j*c + i*c*h + 0*c*h*b + ...
• 线性索引 = k + j*c + i*c*h

窗口和扩展 (Window and Expand)

窗口 (Window)

窗口操作接受一个张量并输出一个子张量。

例如：

b h w c -> b h[10:-10] w[10:-10] c

在这个例子中，我们“裁剪”了一个子张量，其中 h 和 w 维度分别从第10个索引开始，到倒数第10个索引结束，将高度和宽度维度都缩短了20个大小。

扩展 (Expand)

扩展操作接受一个张量并输出一个带有额外填充或边缘扩展的张量。

例如：

b h[10:-10] w[10:-10] c -> b h w c

这将添加填充，使 h 和 w 维度从10个索引开始，到10个索引结束，将高度和宽度维度的大小都增加了20。

现在这要复杂得多，因为你需要额外的信息来填充新索引。

最常见的填充模式是

• 零填充：用零填充新位置（最常用）
• 边缘填充：重复最近的边缘值
• 环绕填充：通过从对边重复来平铺张量
• 反射填充：在边界处镜像张量

填充模式错误是细微bug的常见来源，尤其是在图像处理中，边缘伪影可能会通过模型传播。

数据操作

数据操作则完全是另一回事；张量的实际值会发生修改。这通常是大多数计算工作完成的地方。

映射 (Map)

映射操作接受一个张量并独立地对每个元素执行操作。在机器学习框架中，你可以将其视为对张量应用逐元素函数，如ReLU、tanh或sigmoid。

它是一个函数，接受一个 shape 的张量并输出一个 shape 的张量。

例如：

b h w c -func> b h w c

或者

func(张量<形状>) -> 张量<形状>

缩减 (Reduce)

缩减操作接受一个张量，并使用（通常是关联的）函数，沿着指定轴合并所有子张量。在机器学习框架中，你可以将其视为求和或求积函数。

它是一个函数，接受一个张量并使用（通常是关联的）函数，沿着指定轴合并所有子张量。本质上，它接受一个形状为 `shape` 的张量列表，并输出一个形状为 `shape` 的张量。

例如，在 h 维度上缩减

b h w c -func> b w c

或者

func(列表<张量<形状>>) -> 张量<形状>

转换 (Transform)

转换操作接受一个张量，并独立地对指定轴上的所有子张量执行操作。在机器学习框架中，你可以将其视为对张量的最后一个维度应用softmax，或者独立地对批次中的每张图像应用滤镜。

它是一个函数，接受一个形状为 `shape` 的张量列表，并输出一个形状为 `shape` 的张量列表。

例如：

b h w c / h w c -func> b h w c

或者

func(列表<张量<形状>>) -> 列表<张量<形状>>

Transform 是 map 的广义张量版本。因此，你可以通过 `b h w c / b h w c -func> b h w c` 来重现 map。

扩充 (Amplect)

扩充操作接受多个张量，并对给定对应方向上的每个元素应用标量函数（通常是关联的）。在机器学习框架中，这等同于乘法或加法等广播操作。

Amplect（源自拉丁语 amplector，意为“拥抱”）是我在此处使用的一个非标准名称，以区别于 broadcast-<name>。

例如：

b h w c, c -func> b h w c

或

关联函数(张量<[]>) -> 张量<[]>

视图操作与数据操作的区别

正如你所看到的，视图操作可以一个接一个地执行，只需记录如何索引张量的指令即可。

在 Candle 框架中，执行操作后有时需要调用 contiguous 函数。这是因为重新索引不会改变内存布局，它非常便宜，并且在调用 contiguous 时才进行惰性评估。通常，当你即将执行需要张量在内存中按顺序排列的数据操作时，例如矩阵乘法或将数据传递给需要连续数组的外部库时，你需要调用 contiguous()。

任何给定图中的节点中，出人意料地有很高比例是这些廉价的视图操作。

推导注意力机制

这些并非所有你需要的操作，但它们是其中的大部分。为了向你展示这些操作有多么强大，让我们用它们来创建注意力机制。我将使用 $ 给张量命名，以帮助理清流程。

步骤1

我们从输入张量 $Input{b s d} 开始，其中 b 代表批次大小，s 代表 token 数量，d 代表 token 维度大小；以及 $W_k{d e}，它是我们的 K 权重矩阵，其中 d 代表 token 维度大小，e 代表隐藏层大小。

$Input{b s d}, $W_k{d e} -乘法> $K_p{b s d e} -求和> $K{b s e}

在这一步中，我们可以看到第一个操作是扩充操作，然后是缩减操作，它正在执行矩阵乘法以获得 $K{b s e}。

接着是

$K{b s e} -> $K_T{b e s}

一个用于获取 $K_T{b e s} 的置换操作，其中 e 和 s 维度已交换。

我们对作为输入的价值权重矩阵 $W_v{d e} 遵循类似的过程

$Input{b s d}, $W_v{d e} -乘法> $V_p{b s d e} -求和> $V{b s e}

与上次不同，我们不会置换 $V{b s e}。

同样地，对于 $W_q{d e}，我们得到 $Q{b sq e}

$Input{b s d}, $W_q{d e} -乘法> $Q_p{b s d e} -求和> $Q{b s e}

步骤2

我在 einops 符号中使用 sq 和 sk 而不是仅仅使用 s，以区分查询和键序列维度，尽管它们长度相同。

首先，我们对Q和K^T矩阵进行矩阵乘法，得到注意力矩阵 $A_p{b sq sk}

$Q{b sq e}, $K_T{b e sk} -乘法> $A_pp{b sq e sk} -求和> $A_p{b sq sk}

然后我们对 sk 维度应用 softmax 来计算注意力权重

$A_p{b sq sk}/{sk} -softmax> $A{b sq sk}

然后，我们对最终注意力矩阵和值进行矩阵乘法，以得到输出。

$A{b sq sk}, $V{b sk e} -乘法> $O_p{b sq sk e} -求和> $O{b sq e}

步骤3

使用我们的最后一个输入 $W_o{e d}，我们执行矩阵乘法，将输出维度投射回 token 空间。

$O{b sq e}, $W_o{e d} -乘法> $F_p{b sq e d} -求和> $F{b sq d}

结论

我希望这篇张量操作分类法的介绍对您有所帮助，也许还激发了您对张量操作是多么直接的喜爱。

下次你看到一个重塑操作后面跟着一个矩阵乘法时，你就会知道重塑几乎是免费的，只改变索引指令，而矩阵乘法才是GPU真正进行繁重工作的地方。

这种分类学不仅仅是学术性的。它能帮助你发现优化机会，理解为什么有些操作会成为瓶颈，而有些操作在你的性能分析器中几乎不显示，还能看到像注意力机制这样复杂的算法实际上是如何由这些基本原语组合而成的。

下次你调试性能或设计新的模型架构时，你将能够以廉价的视图操作与昂贵的数据操作来思考。这种心智模型将磨砺你的直觉和张量编程效率。