玩转张量必备利器之 einops
共 4118字，需浏览 9分钟
· 2022-01-10
视觉 Transformer 中免不了要对小批次图像张量以及特征张量一顿操作，比如图像分块、多头自注意力机制等。而这些操作如果借用 einops 以及 einsum 的话往往可以事半功倍。
官方给出的全称是 Einstein-Inspired Notation for operations，也可以看成 Einstein operations。它支持广泛使用的张量包（如 numpy、pytorch、chainer、gluon、tensorflow）。
灵感来自爱因斯坦求和约定（einsum），而这个强大工具可以参考下面这篇。
两大利器 einsum 和 einops 在手，可谓无往不利。
掌握了它们，不仅方便了张量操作，甚至还可以跨包写通用代码。
.安装 .
!pip install einops
1基本套路
在学习 NumPy 的多维数组时，我们知道轴这个概念。对于同一个数组来说，它的几个轴的顺序是可以不同的。
比如，深度学习中经常碰到表示一批次图像数据的张量，往往有不同的表示形式。
再比如，一张图像往往可以表示为一个三维数组，如果将 1 轴和 2 轴交换顺序也是可以的，只是此时同一个元素的索引（下标）变了。
上述变换在 NumPy 中可以通过下面代码来实现，
y = x.transpose(0, 2, 1)
但如果使用 einops，将会是下面这个样子。
y = rearrange(x, 'c h w -> c w h')
形式上明显是借鉴了 einsum，但咋一看，没比 transpose 方便啊。
这个例子是没体现出它的优势，但实际上它还有更多功能，比如多个轴组合、轴分解，分解再组合，约简除轴或者增加新轴等。
2操作单张图片
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = [u'SimHei']
plt.rcParams['axes.unicode_minus'] = False
from utils import display_np_arrays_as_images
display_np_arrays_as_images()
我们加载一幅小扎元宇宙的图片来测试。
img = np.array(Image.open('./resources/meta_verse_256.jpg'))/255
img.shape, img.dtype
((256, 460, 3), dtype('float64'))
img
看这个图的大小，(256, 460, 3)。这三个数字对应高度、宽度、通道数，即 h w c。
from einops import rearrange, reduce, repeat
.元素重排.
einops.rearrange 是一种对多维数组（张量）进行元素重排的十分强大的操作。
该操作包括转置（轴置换）、reshape、挤压（squeeze）、解压（unsqueeze）、堆叠（stack）、拼接（concatenate）等操作。
转置
这个操作很方便，高和宽两个轴交换顺序即可，即 h w c -> w h c。
out = rearrange(img, 'h w c -> w h c')
out.shape, out.dtype
((460, 256, 3), dtype('float64'))
out
out = rearrange(img, 'h w (c cs) -> (h cs) w c', cs=3)
out.shape
(768, 460, 1)
out[...,0]
上面不是将三个通道沿着高度方向拼接，而是将原高度和通道两个轴合并了。是三个通道穿插起来了，表现出来是整个图高度方向被拉长了。
那么怎么做到将三个通道沿着高度方向拼接呢？
先将三个轴重新排序为 c h w格式，一般深度学习的包中会以这种格式处理数据。
out = rearrange(img, 'h w c -> c h w')
out.shape
(3, 256, 460)
然后将上面结果的三个通道沿 h方向拼接，这个目标只需要合并c和h两个轴来实现。
rearrange(out, 'c h w -> (c h) w')
上述代码中，c 轴在 h 轴前面，因此 c 是 0 轴，h 是 1 轴，因此数据总体是按三个通道分开排列。
也可以将上述操作合并成一步
rearrange(img, 'h w c -> (c h) w')
将图片展平为一维数组， 353280 = 256 x 460 x 3。
rearrange(img, 'h w c -> (c h w)').shape
(353280,)
3操作一批图片
上面仅仅是对一张图片进行操作，但像在深度学习中往往是对一批图片下手。
我们不妨也来一试，但不想另外再找图，咱们就地取材，从上面这张图中取小图来操作一番。
将每个图像分成 8 个更小的图像块
bhwc = rearrange(img, '(h1 h) (w1 w) c -> (h1 w1) h w c', h1=2, w1=4)
bhwc.shape
(8, 128, 115, 3)
def subfig(bhwc, hs, ws):
fig, ax = plt.subplots(hs, ws, figsize=(12, 6))
for i, axi in enumerate(ax.flat):
axi.imshow(bhwc[i])
axi.set(xticks=[], yticks=[], xlabel='第 '+str(i+1)+' 个子图')
subfig(bhwc, hs=2, ws=4)
8 个小图块横排
rearrange(bhwc, 'b h w c -> h (b w) c')
8 个小图块纵排
rearrange(bhwc, 'b h w c -> (b h) w c')
bhwc.shape
(8, 128, 115, 3)
空域到通道转换
res = rearrange(img, '(h1 h) (w1 w) c -> (h1 w1) h w c', h1=2, w1=2)
space2deep = rearrange(res, 'b (h1 h) (w1 w) c -> b h w (c h1 w1)', h1=2, w1=2)
space2deep.shape
(4, 64, 115, 12)
上述代码将 4 个子图取出来沿通道轴拼接得到 12 个通道。
将每个图展平
flt = rearrange(bhwc, 'b h w c -> b (h w c)')
flt.shape
(8, 44160)
4Reduce 和 Repeat
einops 也可以作约简（reduce）运算，即沿着某个轴聚合，减少轴；当然也能反过来，沿着新的轴 repeat，增加轴。
对 c轴求均值
out1 = reduce(bhwc, 'b h w c -> b h w', reduction='mean')
out1.shape
(8, 128, 115)
上述代码将每个子图的三个通道作了均值约简，变成单通道灰度图。
下面这样子可以保留被约简的轴。
mean1 = reduce(bhwc, 'b h w c -> b h w 1', 'mean')
mean1.shape
(8, 128, 115, 1)
下面代码将其重新组装成完整的三通道灰度图。
out2 = repeat(out1, '(bh bw) h w -> (bh h) (bw w) c', c=3, bh=2, bw=4)
out2.shape
(256, 460, 3)
out2
reduce 和 repeat 联合使用实现马赛克效果。
res = reduce(img, '(h hs) (w ws) c -> h w c', reduction='mean', hs=8, ws=5)
repeat(res, 'h w c-> (h hs) (w ws) c', c=3, hs=8, ws=5)
.求差 .
min1 = reduce(bhwc, 'b h w c -> b h w ()', 'min')
min1.shape
(8, 128, 115, 1)
dif = bhwc - min1
subfig(dif, hs=2, ws=4)
5加新轴
out3 = rearrange(bhwc, 'b h w c -> 1 b h w 1 c')
out3.shape
(1, 8, 128, 115, 1, 3)
6图像 Patch
将一张图划分为若干个图像子块。
images = rearrange(bhwc, 'b (h ph) (w pw) c -> (b h w) ph pw c', ph = 32, pw = 23)
images.shape
(160, 32, 23, 3)
subfig(images, hs=4, ws=5)
是不是特别适合于拿来实现图像的 Transformer 呢！
7小结
rearrange不改变元素的数量，它的功能涵盖了不同的 numpy 函数（如transpose、reshape、stack、concatenate、squeeze和expand_dims）reduce实现了约简操作（如mean、min、max、sum、prod等）repeat实现重复（repeating）和平铺（tiling）。轴的组合和分解是基石，它们可以并且应该一起使用来操作张量。
相关阅读