einops和einsum：直接操作张量的利器

einops和einsum是Vision Transformer的代码实现里出现的两个操作tensor维度和指定tensor计算的神器，在卷积神经网络里不多见，本文将介绍简单介绍一下这两样工具，方便大家更好地理解Vision Transformer的代码。

einops：直接操作tensor维度的神器

github地址：https://github.com/arogozhnikov/einops

einops：灵活和强大的张量操作，可读性强和可靠性好的代码。支持numpy、pytorch、tensorflow等。

有了他，研究者们可以自如地操作张量的维度，使得研究者们能够简单便捷地实现并验证自己的想法，在Vision Transformer等需要频繁操作张量维度的代码实现里极其有用。

这里简单地介绍几个最常用的函数。

安装

einops的安装非常简单，直接pip即可：

pip install einops

rearrange

import torch
from einops import rearrangei_tensor = torch.randn(16, 3, 224, 224)        # 在CV中很常见的四维tensor： （N，C，H，W）
print(i_tensor.shape)
o_tensor = rearrange(i_tensor, 'n c h w -> n h w c')
print(o_tensor.shape)

输出：

torch.Size([16, 3, 224, 224])
torch.Size([16, 224, 224, 3])

在CV中很常见的四维tensor：（N，C，H，W），即表示（批尺寸，通道数，图像高，图像宽），在Vision Transformer中，经常需要对tensor的维度进行变换操作，rearrange函数可以很方便地、很直观地操作tensor的各个维度。

除此之外，rearrange还有稍微进阶一点的玩法：


i_tensor = torch.randn(16, 3, 224, 224)
o_tensor = rearrange(i_tensor, 'n c h w -> n c (h w)')
print(o_tensor.shape)
o_tensor = rearrange(i_tensor, 'n c (m1 p1) (m2 p2) -> n c m1 p1 m2 p2', p1=16, p2=16)
print(o_tensor.shape)

输出：

torch.Size([16, 3, 50176])
torch.Size([16, 3, 14, 16, 14, 16])

可以进行指定维度的合并和拆分，注意拆分时需要在变换规则后面指定参数。

repeat

from einops import repeati_tensor = torch.randn(3, 224, 224)
print(i_tensor.shape)
o_tensor = repeat(i_tensor, 'c h w -> n c h w', n=16)
print(o_tensor.shape)

repeat时记得指定右侧repeat之后的维度值

输出：

torch.Size([3, 224, 224])
torch.Size([16, 3, 224, 224])

reduce

from einops import reducei_tensor = torch.randn((16, 3, 224, 224))
o_tensor = reduce(i_tensor, 'n c h w -> c h w', 'mean')
print(o_tensor.shape)
o_tensor_ = reduce(i_tensor, 'b c (m1 p1) (m2 p2)  -> b c m1 m2 ', 'mean', p1=16, p2=16)
print(o_tensor_.shape)

输出：

torch.Size([3, 224, 224])
torch.Size([16, 3, 14, 14])

reduce时记得指定左侧要被reduce的维度值

Rearrange

import torch
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
from einops.layers.torch import Rearrangemodel = Sequential(Conv2d(3, 64, kernel_size=3),MaxPool2d(kernel_size=2),Rearrange('b c h w -> b (c h w)'),      # 相当于 flatten 展平的作用Linear(64*15*15, 120), ReLU(),Linear(120, 10)
)i_tensor = torch.randn(16, 3, 32, 32)
o_tensor = model(i_tensor)
print(o_tensor.shape)

输出：

torch.Size([16, 10])

einops.layers.torch.Rearrange 是nn.Module的子类，可以放在网络里面直接当作一层。

torch.einsum：爱因斯坦简记法

爱因斯坦简记法：是一种由爱因斯坦提出的，对向量、矩阵、张量的求和运算 ∑\sum∑ 的求和简记法。

在该简记法当中，省略掉的部分是：

求和符号 ∑\sum∑
求和号的下标 iii

省略规则为：默认成对出现的下标（如下例1中的 iii 和例2中的 kkk ）为求和下标，被省略。

1）xiyix_iy_ixiyi简化表示内积 <x,y><\mathbf{x},\mathbf{y}><x,y>
xiyi:=∑ixiyi=ox_iy_i := \sum_i x_iy_i = o xiyi:=i∑xiyi=o

其中o为输出。

用 XikYkjX_{ik}Y_{kj}XikYkj 简化表示矩阵乘法 XY\mathbf{X}\mathbf{Y}XY
XikYkj:=∑kXikYkj=OijX_{ik}Y_{kj}:=\sum_k X_{ik}Y_{kj}=\mathbf{O}_{ij} XikYkj:=k∑XikYkj=Oij
其中 Oij\mathbf{O}_{ij}Oij 为输出矩阵的第ij个元素。

这样的求和简记法，能够以一种统一的方式表示各种各样的张量运算（内积、外积、转置、点乘、矩阵的迹、其他自定义运算），为不同运算的实现提供了一个统一模型。

einsum在numpy和pytorch中都有实现，下面我们以在torch中为例，展示一下最简单的用法

import torchi_a = torch.randn(16, 32, 4, 8)
i_b = torch.randn(16, 32, 8, 16)out = torch.einsum('b h i j, b h j d -> b h i d', i_a, i_b)
print(out.shape)

输出：

torch.Size([16, 32, 4, 16])

可以看到，torch.einsum可以简便地指定tensor运算，输入的两个tensor维度分别为 bhijb\ h\ i\ jb h i j 和 bhjdb\ h\ j\ db h j d ，经过tensor运算后，得到的张量维度为 bhidb\ h\ i\ db h i d 。代码运行结果与我们的预期一致。