安装

einops是用来操作张量的Python库，其强大之处在于，提供了一套可读性超高的字符串语法，对张量进行重塑、转置、降维、复制等操作。其语法风格与np.einsum十分相似。

安装没什么好说的，pip即可

pip install einops

重排

一张RGB图像，维度是$64\times 64\times 3$，下面将其调整为$3\times 64\times 64$，用einops和numpy实现的方案分别如下，二者差距一目了然

import numpy as np
from einops import rearrange

hwc = np.random.rand(64, 64, 3) 
chw = rearrange(hwc, 'h w c -> c h w')

chwNp = hwc.transpose(2, 0, 1)

对于更复杂的问题，比如将两张$3\times 64\times 64$的图像拆分成多个长度为$3\times 16\times 16$的patch，einops和传统写法的差距将进一步拉大

x = np.random.rand(2, 3, 64, 64)
patches = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)', p1=16, p2=16)
patches.shape   # (2, 16, 768) = (2, ?, 3×16×16)

xNp = x.reshape(2, 3, 64//16, 16, 64//16, 16)
xNp = xNp.transpose(0, 2, 4, 3, 5, 1)
xNp = xNp.reshape(2, -1, 3 * 16 * 16)

比如现有一形如$(b,c,h,w)$的图像序列，希望将其按照patch的尺寸$p$展开为$[b,n,(c,p,p)]$，则用einops

from einops import rearrange
x = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (c p1 p2)', p1=p, p2=p)

降维和重复

下面考虑一个经典的池化问题，对于$2\times 3\times 64\times 64$的数据，对每个通道的数据求平均值或最大值，传统方案和einops的方案对比如下

from einops import reduce

features = np.random.rand(2, 3, 64, 64)

fNp = features.mean(axis=(2, 3))    # 传统写法

pooled = reduce(features, 'b c h w -> b c', 'mean')     # einops平均值
max_pooled = reduce(features, 'b c h w -> b c', 'max')  # einops最大值池化

einops.reduce支持的降维方案包括min, max, sum, mean, prod any, all。

能降就能升，einops中，用repeat可以实现对任意轴进行重复。例如，将$2\times 512$的数据，复制成$2\times 10\times 512$，方法如下

from einops import repeat

emb = np.random.rand(2, 512)
emb_expanded = repeat(emb, 'b d -> b t d', t=10)