tf.matmulとtf.tensordotの違い

はじめに

それぞれnumpy.matmulとnumpy.tensordotと同じです。

いずれも行列をベースに内積を計算する関数ですが、
ブロードキャストへの対応が異なります。

ブロードキャストについては下記など。
https://qiita.com/capybara/items/2d5f387996a5bd275fbf

numpyやtensorflowに慣れた方には常識かもしれませんが、
tf.matmulやtf.tensordotで検索する方のために投稿します。

tf.matmul

a: shape=[x, y, z, i, k]
b: shape=[x, y, z, k, j]

tf.matmul(a,b)
# shape=[x, y, z, i, j]

もし, b: shape=[x,y,z, j, i]なら、

tf.matmul(a,b, transpose_b=True)


shapeからも分かるように、
tf.matmulはブロードキャストしません。
最後の二階部分以外はテンソルの形がそろっていることが必要です。

また、trasnpose_a, transpose_bでは、最後の二階部分のみが転置されます。

tf.tensordot

a: shape=[x, y, z, i, k]
b: shape=[j, k]

tf.tensordot(a,b, axis=[4,1])
# shape=[x, y, z, i, j]

このようにブロードキャストされた内積が計算できます。
なお、axisの指定の仕方はいくつかあって、

• axis=p
  aの最後のp階部分と、bの最初のp階部分についてのdot

• axis=[p, l]
  aのp+1階部分とbのl+1階部分のdot

• axis=[[p,q,r],[l,m,n]]
  aのp+1階部分とbのl+1階部分、
  aのq+1階部分とbのm+1階部分、
  aのr+1階部分とbのn+1階部分、のdot

となっています。

axisで指定されなかった部分に関してbroadcastされ、
最終的なテンソルの形は残った階をa,bの順に並べた形になります。

おまけ

tf.einsumという関数がさらに自由度の高いものとして用意されているようです。。

実例：バッチ処理でマルチラベル分類

最終隠れ層を受け取って、最終的にマルチラベル分類を行うlayerは以下のようになります。
output_layer: shape=[batch_size, hidden_size]
weight: shape=[num_labels, num_classes, hidden_size]
bias: shape[num_labels, num_classes]

logits=tf.tensordot(output_layer, weight, axis=[1,2])
tf.shape(logits) #[batch_size, num_labels, num_classes]

logits=tf.add(logits, bias)
tf.shape(logits) #[batch_size, num_labels, num_classes]

probs=tf.nn.softmax(logits, axis=-1)
tf.shape(probs) #[batch_size, num_tasks, num_labels]

これで、適当な損失関数に渡せます。

参考

公式doc
https://www.tensorflow.org/api_docs/python/tf/linalg/matmul
https://www.tensorflow.org/api_docs/python/tf/tensordot
stackoverflow
https://stackoverflow.com/questions/38051143/no-broadcasting-for-tf-matmul-in-tensorflow/38056381
numpyやchainerでのベクトル、行列、テンソルの積関連演算まとめ
http://segafreder.hatenablog.com/entry/2017/01/09/203811