Einops重排函数基本功能

Question

我正在尝试理解用于张量重新排序的 einops 语法，但我不知何故错过了要点

如果我有以下矩阵：

mat = torch.randint(1, 10, (8,4))

我明白以下命令的作用：

rearrange(mat, '(h n) w -> (n h) w', n = 2)

但我无法真正理解以下内容：

rearrange(mat, '(n h) w -> (h n) w', n = 2)

rearrange(mat, '(n h) w -> (h n) w', n = 4)

如有任何帮助，我们将不胜感激

Answer 1

rearrange(mat, '(h n) w -> (n h) w', n = 2)
and
rearrange(mat, '(n h) w -> (h n) w', n = 2)

互为倒置。如果你能想象一个做什么，第二个做反向变换

至于后者，垫子是8x4

rearrange(mat, '(n h) w -> (h n) w', n = 4)

所以你首先将第一个维度拆分为 4x2（下面我忽略了 w 维度，因为它没有发生任何特殊情况）

[0, 1, 2, 3, 4, 5, 6, 7]

到

[0, 1, 
 2, 3, 
 4, 5, 
 6, 7]

然后将轴的顺序更改为 2x4（转置）

[0, 2, 4, 6,
 1, 3, 5, 7]

然后将两个维度合并为一个

[0, 2, 4, 5, 1, 3, 5, 7]

如果您仍然不明白它是如何工作的，请举一些更简单的例子，比如

rearrange(np.arange(50), '(h n) -> h n', h=5)
rearrange(np.arange(50), '(h n) -> h n', h=10)
rearrange(np.arange(50), '(h n) -> n h', h=10)

等这样您就可以跟踪矩阵中每个元素的移动

Einops rearrange function basic functionallity