tensorflow einsum vs. matmul vs. tensordot

tensorflow einsum vs. matmul vs. tensordot

在tensorflow中,函数tf.einsumtf.matmultf.tensordot都可以用于相同的任务。 (我意识到 tf.einsumtf.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)

有没有比其他人更受欢迎的?

两者都 tf.tensordot() and tf.einsum() are syntactic sugar that wrap one or more invocations of tf.matmul() (although in some special cases tf.einsum() can reduce to the simpler elementwise tf.multiply()).

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