web-dev-qa-db-fra.com

tensorflow einsum vs matmul vs tensordot

Dans tensorflow, les fonctions tf.einsum, tf.matmul Et tf.tensordot Peuvent toutes être utilisées pour les mêmes tâches. (Je me rends compte que tf.einsum Et tf.tensordot Ont des définitions plus générales; je me rends également compte que tf.matmul A une fonctionnalité de traitement par lots.) Dans une situation où l'un des trois pourrait être utilisé, une fonction a tendance à être la plus rapide? Existe-t-il d'autres règles de recommandation?

Par exemple, supposons que A est un tenseur de rang 2 et b est un tenseur de rang 1, et que vous souhaitez calculer le produit c_j = A_ij b_j. Des trois options:

c = tf.einsum('ij,j->i', A, b)

c = tf.matmul(A, tf.expand_dims(b,1))

c = tf.tensordot(A, b, 1)

est-il généralement préférable aux autres?

16
John Kleve

tf.tensordot() et tf.einsum() sont des sucres syntaxiques qui enveloppent une ou plusieurs invocations de tf.matmul() (bien que dans certains cas particuliers, tf.einsum() puisse se réduire au plus simple élément par élément tf.multiply() ).

À la limite, je m'attendrais à ce que les trois fonctions aient des performances équivalentes pour le même calcul. Cependant, pour des matrices plus petites, il peut être plus efficace d'utiliser directement tf.matmul(), car cela produirait un graphique TensorFlow plus simple avec moins d'opérations, et donc les coûts d'invocation par opération seront plus faibles.

13
mrry