tensorflow einsum vs.matmul vs.tensord公司

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.einsum()可以简化为更简单的元素级^{}）。

在极限条件下，我希望这三个函数对于相同的计算都有相同的性能。然而，对于较小的矩阵，直接使用tf.matmul()可能更有效，因为它将产生一个操作较少的更简单的TensorFlow图，因此每次操作的调用成本将更低。