TensorFlow einsum与matmul与tensordot的区别

Question

在Tensorflow中，函数tf.einsum、tf.matmul和tf.tensordot都可以用于相同的任务。（我知道tf.einsum和tf.tensordot有更通用的定义；我也知道tf.matmul具有批处理功能。）在任何一种函数可以使用的情况下，哪个函数更快？还有其他的推荐规则吗？

例如，假设A是一个二阶张量，b是一个一阶张量，你想计算乘积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) 通常有哪个更好呢？

- John Kleve

Answer 1

tf.tensordot() 和tf.einsum() 都是语法糖，用于包装一个或多个tf.matmul()的调用（虽然在一些特殊情况下，tf.einsum()可以简化为更简单的按元素运算tf.multiply()）。

理论上，我期望这三个函数对于相同的计算具有等效的性能。但是，对于较小的矩阵，直接使用tf.matmul()可能更有效，因为它将产生一个带有较少操作的更简单的TensorFlow图，因此预操作调用成本会更低。

- mrry

---

在我的例子中，在应用tf.matmul之前，我必须对b使用tf.expand_dims。此外，使用tf.matmul返回一个秩为2的张量，而不是秩为1的张量；使c成为秩为1的张量需要在矩阵乘法后调用tf.squeeze。这些squeeze和expand_dims操作是否具有有意义的时间成本？ - John Kleve

它们是纯元数据操作，因此具有非常小的恒定成本，应该被 tf.matmul() 本身所主导。 - mrry

你确定吗？从根本上讲，einsum 可能的操作集是 matmul 可能操作集的一个巨大超集；matmul 总是在单个轴上进行收缩。在 numpy 的 einsum 中，可能会存在实质性的性能差异。 - Eelco Hoogendoorn