python - 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 张量，并且您想要计算 product 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()可以简化为更简单的 elementwise tf.multiply()）。

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