一个函数打天下，einsum
极市导读
在实现算法需要转为代码实现的时候遇到复杂的函数，不仅困难还容易出错，本文介绍了einsum函数帮助大家解决这个问题。文章共举例五个函数通过einsum标记法表示，让大家对该函数的运用更加恰当。>>加入极市CV技术交流群，走在计算机视觉的最前沿
einsum全称Einstein summation convention（爱因斯坦求和约定），又称为爱因斯坦标记法，是爱因斯坦1916年提出的一种标记约定，简单的说就是省去求和式中的求和符号，例如下面的公式：
以einsum的写法就是：
后者将 
符号给省去了，显得更加简洁；再比如：
(1)
(2)
上面两个栗子换成einsum的写法就变成：
(1)
(2)
在实现一些算法时，数学表达式已经求出来了，需要将之转换为代码实现，简单的一些还好，有时碰到例如矩阵转置、矩阵乘法、求迹、张量乘法、数组求和等等，若是以分别以transopse、sum、trace、tensordot等函数实现的话，不但复杂，还容易出错
现在，这些问题你统统可以一个函数搞定，没错，就是einsum，einsum函数就是根据上面的标记法实现的一种函数，可以根据给定的表达式进行运算，可以替代但不限于以下函数：
矩阵求迹：trace
求矩阵对角线：diag
张量（沿轴）求和：sum
张量转置：transopose
矩阵乘法：dot
张量乘法：tensordot
向量内积：inner
外积：outer
该函数在numpy、tensorflow、pytorch上都有实现，用法基本一样，定义如下：
einsum(equation, *operands)
equation是字符串的表达式，operands是操作数，是一个元组参数，并不是只能有两个，所以只要是能够通过einsum标记法表示的乘法求和公式，都可以用一个einsum解决，下面以numpy举几个栗子：
# 沿轴计算张量元素之和：
c = a.sum(axis=0)
上面的以sum函数的实现代码，设 
为三维张量，上面代码用公式来表达的话就是：
换成einsum标记法：
然后根据此式使用einsum函数实现等价功能：
c = np.einsum('ijk->jk', a)
# 作用与 c = a.sum(axis=0) 一样
更进一步的，如果 不止是三维，可以将下标 
换成省略号，以表示剩下的所有维度：
c = np.einsum('i...->...', a)
这种写法pytorch与tensorflow同样支持，如果不是很理解的话，可以查看其对应的公式：
# 矩阵乘法
c = np.dot(a, b)
矩阵乘法的公式为：
然后是einsum对应的实现：
c = np.einsum('ij,jk->ik', a, b)
最后再举一个张量乘法栗子：
# 张量乘法
c = np.tensordot(a, b, ([0, 1], [0, 1]))
如果 
是三维的，对应的公式为：
对应的einsum实现：
c = np.einsum('ijk,ijl->kl', a, b)
下面以numpy做一下测试，对比einsum与各种函数的速度，这里使用python内建的timeit模块进行时间测试，先测试（四维）两张量相乘然后求所有元素之和，对应的公式为：
然后是测试代码：
from timeit import Timer
import numpy as np
# 定义两个全局变量
a = np.random.rand(64, 128, 128, 64)
b = np.random.rand(64, 128, 128, 64)
# 定义使用einsum与sum的函数
def einsum():
temp = np.einsum('ijkl,ijkl->', a, b)
def npsum():
temp = (a * b).sum()
# 打印运行时间
print("einsum cost:", Timer("einsum()", "from __main__ import einsum").timeit(20))
print("npsum cost:", Timer("npsum()", "from __main__ import npsum").timeit(20))
上面Timer是timeit模块内的一个类
Timer(stmt, setup).timeit(number)
# stmt: 要测试的语句
# setup: 传入stmt的运行环境，比如stmt中要导入的模块等。
# 可以写一行语句，也可以写多行语句，写多行语句时要用分号；隔开语句
# number: 执行次数
将两个函数各执行20遍，最后的结果为，单位为秒：
einsum cost: 1.5560735
npsum cost: 8.0874927
可以看到，einsum比sum快了几乎一个量级，接下来测试单个张量求和：
将上面的代码改一下：
def einsum():
temp = np.einsum('ijkl->', a)
def npsum():
temp = a.sum()
相应的运行时间为：
einsum cost: 3.2716003
npsum cost: 6.7865246
还是einsum更快，所以哪怕是单个张量求和，numpy上也可以用einsum替代，同样，求均值（mean）、方差（var）、标准差（std）也是一样
接下来测试einsum与dot函数，首先列一下矩阵乘法的公式以以及einsum表达式：
然后是测试代码：
a = np.random.rand(2024, 2024)
b = np.random.rand(2024, 2024)
# einsum与dot比较
def einsum():
res = np.einsum('ik,kj->ij', a, b)
def dot():
res = np.dot(a, b)
print("einsum cost:", Timer("einsum()", "from __main__ import einsum").timeit(20))
print("dot cost:", Timer("dot()", "from __main__ import dot").timeit(20))
# einsum cost: 80.2403851
# dot cost: 2.0842243
这就很尴尬了，比dot慢了40倍（并且差距随着矩阵规模的平方增加），这还怎么打天下？不过在numpy的实现里，einsum是可以进行优化的，去掉不必要的中间结果，减少不必要的转置、变形等等，可以提升很大的性能，将einsum的实现改一下：
def einsum():
res = np.einsum('ik,kj->ij', a, b, optimize=True)
加了一个参数optimize=True，官方文档上该参数是可选参数，接受4个值：
optimize : {False, True, ‘greedy’, ‘optimal’}, optional
optimize默认为False，如果设为True，这默认选择‘greedy（贪心）’方式，再看看速度：
einsum cost: 2.0330937
dot cost: 1.9866218
可以看到，通过优化，虽然还是稍慢一些，但是einsum的速度与dot达到了一个量级；不过numpy官方手册上有个einsum_path，说是可以进一步提升速度，但是我在自己电脑上（i7-9750H）测试效果并不稳定，这里简单的介绍一下该函数的用法为：
path = np.einsum_path('ik,kj->ij', a, b)[0]
np.einsum('ik,kj->ij', a, b, optimize=path)
einsum_path返回一个einsum可使用的优化路径列表，一般使用第一个优化路径；另外，optimize及einsum_path函数只有numpy实现了，tensorflow和pytorch上至少现在没有
最后，再测试einsum与另一个常用的函数tensordot，首先定义两个四维张量的及tensordot函数：
a = np.random.rand(128, 128, 64, 64)
b = np.random.rand(128, 128, 64, 64)
def tensordot():
res = np.tensordot(a, b, ([0, 1], [0, 1]))
该实现对应的公式为：
所以einsum函数的实现为：
def einsum():
res = np.einsum('ijkl,ijmn->klmn', a, b, optimize=True)
tensordot也是链接到BLAS实现的函数，所以不加optimize肯定比不了，最后结果为：
print("einsum cost:", Timer("einsum()", "from __main__ import einsum").timeit(1))
print("tensordot cost:", Timer("tensordot()", "from __main__ import tensordot").timeit(1))
# einsum cost: 4.2361331
# tensordot cost: 4.2580409
def einsum():
temp = einsum('...->', a, optimize=True)
def test():
temp = a.sum()
np.einsum('...->', a, optimize=True) # 正常运行
np.einsum('...->', a) # 报错
np.einsum('i...->...', a) # 正常
np.einsum('...,...->...', a, b) # 正常
np.einsum('i...->i', a, optimize=True) # 必须加optimize，不然报错
c = (a * b).sum()
# 如果不知道a, b的维数，使用einsum实现上面的功能也必须要加optimize
c = einsum('...,...->', a, b, optimize=True)
推荐阅读
