一个函数打天下，einsum

作者丨永远在你身后&＃64;知乎

来源丨https://zhuanlan.zhihu.com/p/71639781

编辑丨极市平台

einsum全称Einstein summation convention&＃xff08;爱因斯坦求和约定&＃xff09;&＃xff0c;又称为爱因斯坦标记法&＃xff0c;是爱因斯坦1916年提出的一种标记约定&＃xff0c;简单的说就是省去求和式中的求和符号&＃xff0c;例如下面的公式&＃xff1a;

以einsum的写法就是&＃xff1a;

后者将  符号给省去了&＃xff0c;显得更加简洁&＃xff1b;再比如&＃xff1a;

 (1)

 (2)

上面两个栗子换成einsum的写法就变成&＃xff1a;

 (1)

 (2)

在实现一些算法时&＃xff0c;数学表达式已经求出来了&＃xff0c;需要将之转换为代码实现&＃xff0c;简单的一些还好&＃xff0c;有时碰到例如矩阵转置、矩阵乘法、求迹、张量乘法、数组求和等等&＃xff0c;若是以分别以transopse、sum、trace、tensordot等函数实现的话&＃xff0c;不但复杂&＃xff0c;还容易出错

现在&＃xff0c;这些问题你统统可以一个函数搞定&＃xff0c;没错&＃xff0c;就是einsum&＃xff0c;einsum函数就是根据上面的标记法实现的一种函数&＃xff0c;可以根据给定的表达式进行运算&＃xff0c;可以替代但不限于以下函数&＃xff1a;

矩阵求迹&＃xff1a;trace
求矩阵对角线&＃xff1a;diag
张量&＃xff08;沿轴&＃xff09;求和&＃xff1a;sum
张量转置&＃xff1a;transopose
矩阵乘法&＃xff1a;dot
张量乘法&＃xff1a;tensordot
向量内积&＃xff1a;inner
外积&＃xff1a;outer

该函数在numpy、tensorflow、pytorch上都有实现&＃xff0c;用法基本一样&＃xff0c;定义如下&＃xff1a;

einsum(equation, *operands)


equation是字符串的表达式&＃xff0c;operands是操作数&＃xff0c;是一个元组参数&＃xff0c;并不是只能有两个&＃xff0c;所以只要是能够通过einsum标记法表示的乘法求和公式&＃xff0c;都可以用一个einsum解决&＃xff0c;下面以numpy举几个栗子&＃xff1a;

# 沿轴计算张量元素之和&＃xff1a;
c &＃61; a.sum(axis&＃61;0)


上面的以sum函数的实现代码&＃xff0c;设  为三维张量&＃xff0c;上面代码用公式来表达的话就是&＃xff1a;

换成einsum标记法&＃xff1a;

然后根据此式使用einsum函数实现等价功能&＃xff1a;

c &＃61; np.einsum(&＃39;ijk->jk&＃39;, a)
# 作用与 c &＃61; a.sum(axis&＃61;0) 一样


更进一步的&＃xff0c;如果  不止是三维&＃xff0c;可以将下标  换成省略号&＃xff0c;以表示剩下的所有维度&＃xff1a;

c &＃61; np.einsum(&＃39;i...->...&＃39;, a)


这种写法pytorch与tensorflow同样支持&＃xff0c;如果不是很理解的话&＃xff0c;可以查看其对应的公式&＃xff1a;

# 矩阵乘法
c &＃61; np.dot(a, b)


矩阵乘法的公式为&＃xff1a;

然后是einsum对应的实现&＃xff1a;

c &＃61; np.einsum(&＃39;ij,jk->ik&＃39;, a, b)


最后再举一个张量乘法栗子&＃xff1a;

# 张量乘法
c &＃61; np.tensordot(a, b, ([0, 1], [0, 1]))


如果  是三维的&＃xff0c;对应的公式为&＃xff1a;

对应的einsum实现&＃xff1a;

c &＃61; np.einsum(&＃39;ijk,ijl->kl&＃39;, a, b)


下面以numpy做一下测试&＃xff0c;对比einsum与各种函数的速度&＃xff0c;这里使用python内建的timeit模块进行时间测试&＃xff0c;先测试&＃xff08;四维&＃xff09;两张量相乘然后求所有元素之和&＃xff0c;对应的公式为&＃xff1a;

然后是测试代码&＃xff1a;

from timeit import Timer
import numpy as np# 定义两个全局变量
a &＃61; np.random.rand(64, 128, 128, 64)
b &＃61; np.random.rand(64, 128, 128, 64)# 定义使用einsum与sum的函数
def einsum():temp &＃61; np.einsum(&＃39;ijkl,ijkl->&＃39;, a, b)def npsum():temp &＃61; (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的运行环境&＃xff0c;比如stmt中要导入的模块等。# 可以写一行语句&＃xff0c;也可以写多行语句&＃xff0c;写多行语句时要用分号&＃xff1b;隔开语句# number: 执行次数


将两个函数各执行20遍&＃xff0c;最后的结果为&＃xff0c;单位为秒&＃xff1a;

einsum cost: 1.5560735
npsum cost: 8.0874927


可以看到&＃xff0c;einsum比sum快了几乎一个量级&＃xff0c;接下来测试单个张量求和&＃xff1a;

将上面的代码改一下&＃xff1a;

def einsum():temp &＃61; np.einsum(&＃39;ijkl->&＃39;, a)def npsum():temp &＃61; a.sum()


相应的运行时间为&＃xff1a;

einsum cost: 3.2716003
npsum cost: 6.7865246


还是einsum更快&＃xff0c;所以哪怕是单个张量求和&＃xff0c;numpy上也可以用einsum替代&＃xff0c;同样&＃xff0c;求均值&＃xff08;mean&＃xff09;、方差&＃xff08;var&＃xff09;、标准差&＃xff08;std&＃xff09;也是一样

接下来测试einsum与dot函数&＃xff0c;首先列一下矩阵乘法的公式以以及einsum表达式&＃xff1a;

然后是测试代码&＃xff1a;

a &＃61; np.random.rand(2024, 2024)
b &＃61; np.random.rand(2024, 2024)# einsum与dot比较
def einsum():res &＃61; np.einsum(&＃39;ik,kj->ij&＃39;, a, b)def dot():res &＃61; 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


这就很尴尬了&＃xff0c;比dot慢了40倍&＃xff08;并且差距随着矩阵规模的平方增加&＃xff09;&＃xff0c;这还怎么打天下&＃xff1f;不过在numpy的实现里&＃xff0c;einsum是可以进行优化的&＃xff0c;去掉不必要的中间结果&＃xff0c;减少不必要的转置、变形等等&＃xff0c;可以提升很大的性能&＃xff0c;将einsum的实现改一下&＃xff1a;

def einsum():res &＃61; np.einsum(&＃39;ik,kj->ij&＃39;, a, b, optimize&＃61;True)


加了一个参数optimize&＃61;True&＃xff0c;官方文档上该参数是可选参数&＃xff0c;接受4个值&＃xff1a;

optimize : {False, True, ‘greedy’, ‘optimal’}, optional


optimize默认为False&＃xff0c;如果设为True&＃xff0c;这默认选择‘greedy&＃xff08;贪心&＃xff09;’方式&＃xff0c;再看看速度&＃xff1a;

einsum cost: 2.0330937
dot cost: 1.9866218


可以看到&＃xff0c;通过优化&＃xff0c;虽然还是稍慢一些&＃xff0c;但是einsum的速度与dot达到了一个量级&＃xff1b;不过numpy官方手册上有个einsum_path&＃xff0c;说是可以进一步提升速度&＃xff0c;但是我在自己电脑上&＃xff08;i7-9750H&＃xff09;测试效果并不稳定&＃xff0c;这里简单的介绍一下该函数的用法为&＃xff1a;

path &＃61; np.einsum_path(&＃39;ik,kj->ij&＃39;, a, b)[0]
np.einsum(&＃39;ik,kj->ij&＃39;, a, b, optimize&＃61;path)


einsum_path返回一个einsum可使用的优化路径列表&＃xff0c;一般使用第一个优化路径&＃xff1b;另外&＃xff0c;optimize及einsum_path函数只有numpy实现了&＃xff0c;tensorflow和pytorch上至少现在没有

最后&＃xff0c;再测试einsum与另一个常用的函数tensordot&＃xff0c;首先定义两个四维张量的及tensordot函数&＃xff1a;

a &＃61; np.random.rand(128, 128, 64, 64)
b &＃61; np.random.rand(128, 128, 64, 64)def tensordot():res &＃61; np.tensordot(a, b, ([0, 1], [0, 1]))


该实现对应的公式为&＃xff1a;

所以einsum函数的实现为&＃xff1a;

def einsum():res &＃61; np.einsum(&＃39;ijkl,ijmn->klmn&＃39;, a, b, optimize&＃61;True)


tensordot也是链接到BLAS实现的函数&＃xff0c;所以不加optimize肯定比不了&＃xff0c;最后结果为&＃xff1a;

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


测试了10多次&＃xff0c;基本上速度一样&＃xff0c;einsum表现好一点的&＃xff1b;不过说是一个函数打天下&＃xff0c;肯定是做不到的&＃xff0c;还有一些数组的分割、合并、指数、对数等功能没法实现&＃xff0c;需要使用别的函数&＃xff0c;其他的基本都可以用einsum来实现&＃xff0c;简单而又高效

之后经过进一步测试发现&＃xff0c;优化反而出现速度降低的情况&＃xff0c;例如&＃xff1a;

def einsum():temp &＃61; einsum(&＃39;...->&＃39;, a, optimize&＃61;True)def test():temp &＃61; a.sum()


上面两中对数组求和的方法&＃xff0c;当a是一维向量时&＃xff0c;或者a是多维但是规模很小是&＃xff0c;优化的einsum反而更慢&＃xff0c;但是去掉optimize参数后表现比内置的sum函数稍好&＃xff0c;我认为优化是有一个固定的成本

还有一个坑需要注意的是&＃xff0c;有些情况的省略号不加optimize会报错&＃xff0c;就拿上面的栗子而言&＃xff1a;

np.einsum(&＃39;...->&＃39;, a, optimize&＃61;True) # 正常运行
np.einsum(&＃39;...->&＃39;, a) # 报错


很无奈&＃xff0c;试了很多次&＃xff0c;不加optimize就是会报错&＃xff0c;但是并不是所有的省略号写法都需要加optimize&＃xff0c;例如&＃xff1a;

使用省略号实现上面两个公式并不需要加optimize&＃xff0c;能够正常运行

np.einsum(&＃39;i...->...&＃39;, a) # 正常
np.einsum(&＃39;...,...->...&＃39;, a, b) # 正常


但是如果碰到下面的公式&＃xff1a;

上式表示将a除第一个维度之外&＃xff0c;剩下的维度全部累加&＃xff0c;这种实现就必须要加optimize

np.einsum(&＃39;i...->i&＃39;, a, optimize&＃61;True) # 必须加optimize&＃xff0c;不然报错


再举一个栗子&＃xff1a;

c &＃61; (a * b).sum()
# 如果不知道a, b的维数&＃xff0c;使用einsum实现上面的功能也必须要加optimize
c &＃61; einsum(&＃39;...,...->&＃39;, a, b, optimize&＃61;True)


总结一下&＃xff0c;在计算量很小时&＃xff0c;优化因为有一定的成本&＃xff0c;所以速度会慢一些&＃xff1b;但是&＃xff0c;既然计算量小&＃xff0c;慢一点又怎样呢&＃xff0c;而且使用优化之后&＃xff0c;可以更加肆意的使用省略号写表达式&＃xff0c;变量的维数也不用考虑了&＃xff0c;所以建议无脑使用优化。

觉得有用麻烦给个在看啦~