einops的rearrange函数

einops的rearrange函数

在einops库中，rearrange函数是一个非常强大的工具，用于对张量进行重排和变换。它允许用户以一种直观的方式指定张量的形状变化，而无需手动编写复杂的索引操作。

import torch
from einops import rearrange

rearrange(Tensor, "输入模式 -> 输出模式", **参数)

其中，输入模式和输出模式都是分别给Tensor的维度取名，然后通过箭头->指定如何重新排列这些维度。参数可以用于指定维度的大小或其他属性。

几个例子：

x = torch.arange(6).reshape(2, 3)
# tensor([[0, 1, 2],
#         [3, 4, 5]])

y = rearrange(x, "h w -> w h")
# tensor([[0, 3],
#         [1, 4],
#         [2, 5]])

# shape: (2, 3) → (3, 2)

"h w -> w h"，这段的意思是：

• 输入张量的第一个维度命名为h，第二个维度命名为w。即此时，h=2，w=3
• 在输出张量中，我们交换了这两个维度的位置，输出张量的第一个维度是w，第二个维度是h。因此，输出张量的形状为(3, 2)
• 也就是输出张量中\(x_{ji}\)对应于输入张量中的\(x_{ij}\)

几个要求

• 输入模式和输出模式中，维度名称的数量必须相同，或者说，输入模式中每个维度名称在输出模式中必须出现一次。
• 输出模式中不可以有新的维度名称，必须是输入模式中已有的维度名称。

再来一个例子：

x = torch.arange(6)
# x = [0, 1, 2, 3, 4, 5]

y = rearrange(x, "(h w) -> h w", h=2)
# tensor([[0, 1, 2],
#         [3, 4, 5]])
z = rearrange(x, "(h w) -> w h", h=2)
# tensor([[0, 3],
#         [1, 4],
#         [2, 5]])

在上面的例子中，(h w)虽然使用了两个字母，但是在这里，一个()表示的是一个维度，其大小为括号内所有维度的乘积。也就是说，(h w)表示一个大小为h*w的一维张量。我们用这种方式去期待这个一维张量实际上是由h * w的二维张量展开得到的。而在这里我们输入的张量x是一个一维张量，其大小为6，因此此时有

\[h * w = 6 \]

在输出中，我们指定了输出的形状为h w，并且通过参数h=2指定了h的大小为2，这也意味着程序会自动计算出w=3，从而将输入的一维张量重排为一个(2, 3)的二维张量，当然元素的排列顺序是按照原始张量的顺序进行的。

在计算z的时候，输出的形状为w h，此时经过的流程如下：

(h w) -> h w -> w h

因此这相当于

z = rearrange(y, "h w -> w h")

不妨再来看一个逆过程的例子：

x = torch.arange(6).reshape(2, 3)
# tensor([[0, 1, 2],
#         [3, 4, 5]]) (2, 3)

y = rearrange(x, "h w -> (h w)")
# tensor([0, 1, 2, 3, 4, 5])

z = rearrange(x, "h w -> (w h)")
# tensor([0, 3, 1, 4, 2, 5])
# 2 3 -> 3 2 -> (3 2) = 6

# 两个结果的 shape 都是 (6,)

上面对于y的计算很trivial，而对于z的计算，依旧可以视作：

z_ = rearrange(x, "h w -> w h")
z  = rearrange(z_, "h w -> (h w)")

注意，表达式中维度的名称并不只是代表维度大小，输入输出的维度名称实际上是绑定的。

x = torch.randn(2, 3, 8)

y = rearrange(x, "b s (h d) -> b h s d", h=2)

# shape: (2, 3, 8) → (2, 2, 3, 4)

在上面的代码中，由前面的介绍，不难知道有(h d) = 8，因此程序会自动计算出d=4，从而将输入的张量重排为一个(2, 2, 3, 4)的四维张量。

比如：

x = torch.tensor([[
    [0, 1, 2, 3],   # token0
    [4, 5, 6, 7],   # token1
]])
# shape: (b, s, h*d) = (1, 2, 4)

y = rearrange(x, "b s (h d) -> b h s d", h=2)
# shape: (1, 2, 2, 2)

我们分析一下这个例子：

z = rearrange(x, "b s (h d) -> b s h d", h=2)

# z = tensor([[[[0, 1],[2, 3]], 
#             [[4, 5],[6, 7]]]])

# 下面的就可能有些抽象了，不过我们可以这样：忽略最最后一维，让其等于一个东西
# 即 z = tensor([[[z0,z1], 
#                [z2,z3]]])
#   z_i (2,)

y = rearrange(z, "b s h d -> b h s d")
# 这相当于对z的s 与 h 维上进行转置操作，即
# y = tensor([[z0, z2],
#             [z1, z3]])
# 展开即为
# y = tensor([[[[0, 1],[4, 5]],
#             [[2, 3],[6, 7]]]])

当然，如果我们并不关心前置维度，如batch，我们可以使用...来表示任意数量的前导维度，这对于处理批量数据或高维张量非常有用。...可以匹配任意数量的维度，使得我们可以在不明确指定所有维度的情况下进行操作。

x = torch.randn(2, 4, 3, 6)

y = rearrange(x, "... (k p) -> ... k p", p=2)

# shape: (2, 4, 3, 6) → (2, 4, 3, 3, 2)