Тензорный поток einsum против matmul против tenordot

В тензорном потоке функции 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)

является вообще предпочтительным для других?

Ответы

Ответ 1

Оба tf.tensordot() и tf.einsum() являются синтаксический сахар, который обертывает одно или несколько вызовов tf.matmul() (хотя в некоторых особых случаях tf.einsum() может уменьшаться до более простого элементарного tf.multiply()).

В пределе я ожидал бы, что все три функции будут иметь эквивалентную производительность для одного и того же вычисления. Однако для меньших матриц может быть более эффективно использовать tf.matmul() напрямую, потому что это даст более простой график TensorFlow с меньшим количеством операций, и, следовательно, затраты на вызовы для каждой операции будут ниже.

Ответ 2

Скорость tf.einsum зависит от оптимизации пакетом opt_einsum, который был включен в numpy 1.14 (выпущен в январе 2018 года).

Оператор @использует другой путь кода: http://github.com/tensorflow/tensorflow/issues/1062

Matmul не имеет такой оптимизации.