原创

tensorflow 乘法最强利器: einsum

版权声明:本文为博主原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。
本文链接:https://zhuyulab.blog.csdn.net/article/details/87869527

对应的问题:

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

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

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

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

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

也就是说,对每一个样本,对矩阵Y (A×BA\times B)点乘一个标量 x\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更快, 但是作为例子说明是最容易入门的:
转置可以表示为如下操作:

Bji=Aij\mathbf{B}_{ji} = \mathbf{A}_{ij}

他使用einsum 记法表示为: ij->ji,意思也很明确, 就是->的左边代表原来的矩阵A\mathbf{A}, 右边代表我想生成的矩阵B\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=ΣiΣjAij c = \Sigma_i\Sigma_j\mathbf{A}_{ij}
    对应于代码:
c = tf.einsum('ij->', A)
  1. 同理, 列求和和行求和分别对应于
c = tf.einsum('ij->j', A) #列求和
c = tf.einsum('ij->i', A) #行求和

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

3. 带batch的计算

这里, 通过enisum实现tf.matmul的功能:
如:
Cijl=ΣkAijkBikl C_{ijl} = \Sigma_kA_{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到要做一个矩阵乘法。

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

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

也就是说,对每一个样本,对矩阵Y (A×BA\times B)点乘一个标量 x\mathbf{x}

那么,使用einsum我们就可以表示为:

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

实在是太方便了。

总结

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

2019.2.21 LT

文章最后发布于: 2019-02-21 22:10:25
展开阅读全文
0 个人打赏
私信求帮助
曾参与过风云系列卫星、碳卫星、子午工程、嫦娥等项目的数据处理工作;有超10年大型项目的开发经验。 专栏收入了作者为Python爱好者精心打造的多篇文章,从小白入门学习的基础语法、基础模块精讲等内容外,还提出了“Python语感训练”的概念和方法,不仅为初学者提供了进阶之路,有一定基础的程序员亦可从中受益。后续,本专栏还将加入2D/3D应用开发、数据处理、实战项目等精品内容,敬请期待。
xufive ¥1.90 987人订阅

没有更多推荐了,返回首页

©️2019 CSDN 皮肤主题: 大白 设计师: CSDN官方博客

分享到微信朋友圈

×

扫一扫,手机浏览