tensorflow 乘法最强利器： einsum

对应的问题：

在用tensorflow构造自己的损失函数时，经常会涉及到复杂的矩阵乘法。而这些矩阵乘法本来并不复杂， 比如只是简单的 维度为$A\times B$ 的矩阵 $\mathbf{X}$ 和 维度为 $B\times C$的矩阵$\mathbf{Y}$相乘。 但是由于在深度学习中需要利用并行计算来加快速度，因此在包括keras等框架下，都会多出一维batch维度。 这极大地影响了损失函数的实现， 因为在你自定义的loss function中， $\mathbf{X}$的维度为 $N\times A\times B$， $\mathbf{Y}$的维度为 $N\times B\times C$，这种情况下矩阵相乘就不像在没有batch维度时那么简单了。

事实上， 其实我们做一个for循环，对每一个样本进行 $\mathbf{X}\times \mathbf{Y}$ 的计算就可以了， 这也非常容易理解。 但是很显然， 首先这样的一种操作必须启用动态图机制， 其次， 未经优化的for循环很可能成为整个训练时间加快的瓶颈所在。

另一点就是很可能造成的冗余计算。 比如 $\mathbf{x}$ 为 $1\times A$的向量， $\mathbf{y}$ 为 $A\times 1$的向量, 假设loss就是$\mathbf{x}\times \mathbf{y}$。 由于传入到loss function中的是N (batch_size)个样本， 因此实际上你拿到的 $\mathbf{x}$ 为 $N\times A$， $\mathbf{y}$ 为 $A\times N$ （需要手动转置一下）。 如果直接使用tf.matmul计算，会得到一个$N\times N$ 的矩阵，但事实上， 你只需要N个对角线元素， 也就是N个样本各自的$\mathbf{x}\times \mathbf{y}$ 的结果，非对角线元素的计算无疑是冗余的。 解决这一问题， keras下有K.backend.batch_dot(x,y), 可以只计算对角线元素，这一问题得以顺利解决。 同理， tensorflow下也考虑到了矩阵相乘的情况，tf.matmul(A,B)在A 和 B第一维都是batch维的情况下会自动保持第一维不变，对剩下两维进行正常的矩阵运算，也就是我们想用for循环实现的效果。

然而我们还会碰到很多不同的情况， 比如 $\mathbf{x}$ 为 $N\times 1$的向量， $\mathbf{Y}$ 为 $N\times A\times B$的三维矩阵。 我们要实现这样一个操作

	for i in range(N):
		temp(i, :, :) = Y[i, :, :] .* x(i)

也就是说，对每一个样本，对矩阵Y （$A\times B$）点乘一个标量 $\mathbf{x}$。 而这样的操作很难找到直接对应的可供实现的tf API了。 比如当你使用tf.matmul(x, Y)，会提示你维度无法对上，因为你要实现的是点乘操作，而matmul针对的是正常的矩阵乘法。

还有种种这样的问题， 一个在matlab或者numpy种轻描淡写就可以解决的矩阵乘法问题， 在tensorflow中因为多出的一维batch，导致寸步难行。 但这一切，随着einsum的使用，可以迎刃而解。

einsum的作用：

einsum实现了大一统。 具体来说是这样， 所有上述所说的问题，都可以表示为 einsum记法，从而统一使用einsum API处理。 也就是说，只需改变少量的参数， 就可以应对不同的令人郁闷的乘法问题。

具体使用：

==从这里开始， 步入正题：
如何使用einsum API 来实现上述需求的乘法？

很简单， 只需要一步， 将需要完成的运算用einsum记法表示出来：

根据我多年的学习经验， 如同玩游戏一样， 与其写一堆看似详尽实则令人懵圈的游戏说明，还不如直接开几把示范局让大家自己动手更有体会。 因此， 下面直接通过几个生动而实用的例子，来起到举一反三直接掌握einsum的作用：

1. 转置操作：

这是最简单的操作，虽然我感觉直接使用tf. transpose更快， 但是作为例子说明是最容易入门的：
转置可以表示为如下操作:

$${\mathbf{B}}_{ji}={\mathbf{A}}_{ij}$$

他使用einsum 记法表示为： ij->ji，意思也很明确， 就是->的左边代表原来的矩阵$\mathbf{A}$, 右边代表我想生成的矩阵$\mathbf{B}$。 注意，这里的ij并非固定的必须这两个字母，事实上任何没有歧义的字母都可以代替， 比如 ab->ba，可以完成完全一样的结果。 ij只是作为一种示例而已。
那么在tf中如何实现呢，代码如下：

B= tf.einsum('ab->ba', A) 
#B= tf.einsum('ij->ji', A)  #结果一样

这里A是一个二维矩阵。 可以看到，在tf中的实现非常简单，第一个参数为einsum方程记法，后续只需把需要用到的矩阵依次输入即可，在后面的例子中会有更明显的体现。

2. 求和

学会了这个操作，就可以代替掉tf中的reduce_sum等一系列API了， 所以说enisum一个API统一了无数API的功能。

1. 总求和
   $$c={\Sigma }_i{\Sigma }_j{\mathbf{A}}_{ij}$$
   对应于代码：

c = tf.einsum('ij->', A)

2. 同理， 列求和和行求和分别对应于

c = tf.einsum('ij->j', A) #列求和
c = tf.einsum('ij->i', A) #行求和

虽然我们提供的参数非常少， 但是确实，通过einsum记法，我们准确向程序转述了我们的乘法需要 （通过下标明确了）。

3. 带batch的计算

这里， 通过enisum实现tf.matmul的功能：
如：
$$C_{ijl}={\Sigma }_k{A}_{ijk}{B}_{ikl}$$
即可轻松的表示为

C = tf.einsum('ijk,ikl->ijl', A, B) 

可以看到， 和之前只有一个矩阵的情况不同。 这里->前面要依次输入两个矩阵在einsum中的下标， 然后A,B也要依次作为参数传入。但是理解的话就觉得， 如果不是这样设计，反而奇怪呢。
在这个例子中， 虽然我们只传入了'ijk,ikl->ijl'，就让系统get到了我们的想要的结果。 真的是非常便利的，这也是enisum的强大之处，因此必须安利一波。 同时我们可以看到，这样的一个表达是没有歧义的，首先下标i不变， 系统就会知道i维度不进行乘法操作。 而后面两维中jk,kl->jl， 让系统轻松get到要做一个矩阵乘法。

学习到这个之后，我也解决了在一开始提到的那个问题：
比如 $\mathbf{x}$ 为 $N\times 1$的向量， $\mathbf{Y}$ 为 $N\times A\times B$的三维矩阵。 我们要实现这样一个操作

	for i in range(N):
		temp(i, :, :) = Y[i, :, :] .* x(i)

也就是说，对每一个样本，对矩阵Y （$A\times B$）点乘一个标量 $\mathbf{x}$。

那么，使用einsum我们就可以表示为：

temp = tf.einsum('ijk, i->ijk', Y, x)

实在是太方便了。

总结

einsum除了上手时需要对这个einsum记法有所了解时比较生涩，但上手之后就会体会到其强大之处。 用一个统一的API来进行一切你所能想象到的矩阵乘法，真的是tensorflow中不可多得的语法糖了。因此，值得一篇博客专门来记录，希望方便后来的读者，也为大家做一波安利。

2019.2.21 LT