所有产品

einsum一个函数走天下

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

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

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

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

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

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

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

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

  可以看到,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 与另一个常用的函数 tensordot,首先定义两个四维张量的及 tensordot 函数:

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

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

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

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

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

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

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