einsum,一个函数走天下

作者 |永远在你身后

转载自知乎

【导读】einsum 全称 Einstein summation convention(爱因斯坦求和约定),又称为爱因斯坦标记法,是爱因斯坦 1916 年提出的一种标记约定,本文主要介绍了einsum 的应用。

简单的说,应用einsum就是省去求和式中的求和符号,例如下面的公式:

以 einsum 的写法就是:

后者将符号给省去了,显得更加简洁;再比如:

上面两个栗子换成 einsum 的写法就变成:

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

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

矩阵求迹:trace

求矩阵对角线:diag

张量(沿轴)求和:sum

张量转置:transopose

矩阵乘法:dot

张量乘法:tensordot

向量内积:inner

外积:outer

该函数在 numpy、tensorflow、pytorch 上都有实现,用法基本一样,定义如下:

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

上面的以 sum 函数的实现代码,设为三维张量,上面代码用公式来表达的话就是:

换成 einsum 标记法:

然后根据此式使用 einsum 函数实现等价功能:

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

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

矩阵乘法的公式为:

然后是 einsum 对应的实现:

最后再举一个张量乘法栗子:

如果是三维的,对应的公式为:

对应的 einsum 实现:

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

然后是测试代码:

上面 Timer 是 timeit 模块内的一个类

将两个函数各执行 20 遍,最后的结果为,单位为秒:

可以看到,einsum 比 sum 快了几乎一个量级,接下来测试单个张量求和:

将上面的代码改一下:

相应的运行时间为:

还是 einsum 更快,所以哪怕是单个张量求和,numpy 上也可以用 einsum 替代,同样,求均值(mean)、方差(var)、标准差(std)也是一样。

接下来测试 einsum 与 dot 函数,首先列一下矩阵乘法的公式以以及 einsum表达式:

然后是测试代码:

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

加了一个参数 optimize=True,官方文档上该参数是可选参数,接受4个值:

optimize 默认为 False,如果设为 True,这默认选择‘greedy(贪心)’方式,再看看速度:

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

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

最后,再测试 einsum 与另一个常用的函数 tensordot,首先定义两个四维张量的及 tensordot 函数:

该实现对应的公式为:

所以 einsum 函数的实现为:

tensordot 也是链接到 BLAS 实现的函数,所以不加 optimize 肯定比不了,最后结果为:

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

经过进一步测试发现,优化反而出现速度降低的情况,例如:

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

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

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

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

但是如果碰到下面的公式:

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

再举一个栗子:

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

https://zhuanlan.zhihu.com/p/71639781

(*本文为AI科技大本营转载文章,转载请联系作者)

福利时刻

  • 发表于:
  • 原文链接https://kuaibao.qq.com/s/20190904A0PM4400?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。

扫码关注云+社区

领取腾讯云代金券

年度创作总结 领取年终奖励