tensorflow einsum vs matmul vs tensordot

Question

En 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 réalise également que tf.matmul a la fonctionnalité par lots.) Dans une situation où l'un des trois pourrait être utilisé, est-ce qu'une fonction a tendance à être la plus rapide? Y a-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. Sur les 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 tout préférable généralement aux autres?

Answer 1

Les deux tf.tensordot() et tf.einsum() sont le sucre syntaxique qui enveloppent une ou plusieurs invocations de tf.matmul() (bien que dans certains cas particuliers tf.einsum() peut réduire à la simple tf.multiply() par éléments).

Dans 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 tf.matmul() directement, car cela donnerait un graphique TensorFlow plus simple avec moins d'opérations, et par conséquent les coûts d'invocation par opération seront plus faibles.