tensorflow einsum vs. matmul vs. tensordot

Question

在tensorflow中,函数tf.einsum,tf.matmul和tf.tensordot都可以用于相同的任务.(我意识到tf.einsum并且tf.tensordot有更多的一般定义;我也意识到它tf.matmul具有批处理功能.)在可以使用三者中的任何一种的情况下,一个函数是否最快？还有其他推荐规则吗？

例如,假设这A是一个rank-2张量,并且b是rank-1张量,并且您想要计算产品c_j = A_ij b_j.在三个选项中:

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

c = tf.matmul(A, tf.expand_dims(b,1))

c = tf.tensordot(A, b, 1)

是否通常比其他人更好？

Answer 1

这两个tf.tensordot()和tf.einsum()是包装的一个或多个调用语法糖tf.matmul()(虽然在一些特殊情况下tf.einsum()可以降低到简单的按元素tf.multiply()).

在极限中,我希望所有三个函数在同一计算中具有相同的性能.但是,对于较小的矩阵,tf.matmul()直接使用可能更有效,因为它会产生更简单的TensorFlow图,操作更少,因此每操作调用成本会更低.