tensorflow einsum vs. matmul vs. tensordot

Question

在tensorflow中，函数tf.einsum、tf.matmul和tf.tensordot都可以用于相同的任务。(我意识到tf.einsum和tf.tensordot有更通用的定义；我也意识到tf.matmul具有批处理功能。)在这三个函数中的任何一个都可以使用的情况下，是否有一个函数往往是最快的？还有其他推荐规则吗？

例如，假设A是秩2张量，b是秩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图，因此每个操作的调用成本将更低。