テンソル同士の行列積を計算するにはnp.dot？np.matmul？np.tensordot？
テンソルと行列、テンソルとテンソルの積について、どの使えばいいのか（np.dot, np.matmul, np.tensordot）わからなくなることがあります。アフィン変換の例を通じてどの関数を使えばいいのか見ていきます。
目次
背景
アフィン変換同士の合成は行列積で表されます。しかし、複数のアフィン変換の行列（アフィン変換のテンソル）を、同時に合成したいときはどの関数を使って合成すればよいのでしょうか？
アフィン変換
アフィン変換の行列について振り返ると、xy座標におけるアフィン変換の行列は3×3行列で表されます。これは変換が1個しかないようなケースです。
例えば、画像をx軸方向にa倍, y軸方向にb倍拡大する処理は、
$$A_1 = \begin{bmatrix} a & 0 & 0 \\ 0 & b & 0 \\ 0 & 0 & 1\end{bmatrix}$$
という行列で表されます。この変換を任意の点の集合（3xN 行列）
$$P=\begin{bmatrix} x_1 & x_2 & \cdots & x_N \\ y_0 & y_1 & \cdots & y_N \\ 1 & 1 & \cdots & 1 \end{bmatrix}$$
というに対して行うときは、変換後の点は、
$$P’=A_1P$$
という行列積を取れば良いです。P’は3xN行列になります。変換後のx,y座標を見たいときは、3行目の1の羅列を無視してください。この1は行列積を取る都合上入れているものです。
行列の概念を使うことで、複数の点に対して同一の変換を一気にかけることができるのがアフィン変換の強みです（これはGPUを使った処理では非常に有利に機能します）。
アフィン変換の合成
アフィン変換は拡大・縮小だけではありません。平行移動、せん断、回転もできます。x方向にc、y方向にdだけ平行移動したいときは、
$$A_2 = \begin{bmatrix} 1 & 0 & c \\ 0 & 1 & d \\ 0 & 0 & 1\end{bmatrix}$$
という変換の行列になります。変換後の点の求め方は拡大のケースと同じです。
アフィン変換のすごいところは、複数の変換の合成を行列積で簡単にできるということです（詳細は割愛しますが逆変換も逆行列を取るだけでいいのでとても楽です）。例えば$A_1$と$A_2$の合成変換は、
$$A=A_1A_2$$
という行列積を取ると、合成されたアフィン変換の行列になります。
もしこれをアフィン変換に使わないなら、「x方向に2倍拡大して、横に50ピクセル移動して……あれ、何やってたんだっけ？」とバグの温床になりがちです。つまり、コードを簡潔にするためにも行列計算一本で統一できるのならそちらのほうが良いというメリットがあります。
複数のアフィン変換を同時に計算できるように拡張
今アフィン変換の行列が1種類でしたが、アフィン変換とはただの行列積なので、複数のアフィン変換を複数の点に対して同時に行って計算するということができます。
複数のアフィン変換とはどういうことかというと、shapeは何でもよいですが、x方向に5個、y方向に5個異なるアフィン変換があったとします。つまり、3×3の行列が5×5個あるということになります。全体のアフィン変換のテンソル（行列）を「$T_1$:(5, 5, 3, 3)」というshapeにします。
点の集合Pは「(3, N)」という行列のままです。
また、アフィン変換のテンソル同士の合成をしたいので、$T_1$と同じshapeの「$T_2$：(5, 5, 3, 3)」というテンソルを用意します。
ここで問題です。
- (5, 5, 3, 3)と(5, 5, 3, 3)のshapeのテンソル同士の積を取り、(5, 5, 3, 3)というテンソルを返すにはどういう関数を使えばよいでしょう？
- (5, 5, 3, 3)と(3, N)のshapeのテンソル同士の積を取り、(5, 5, 3, N)というテンソルを返すにはどういう関数を使えばよいでしょう？
一般化して考える
アフィン変換を離れて一般的な4階テンソル同士、テンソルと行列の積の問題として考えてみます。1, 2は、
- (a, b, c, d)と(a, b, d, e)の積を取り、(a, b, c, e)を返すには？
- (a, b, c, d)と(d, N)の積を取り、(a, b, c, N)を返すには？
愚直にforループで書くと次のようになります。
# t1 : (a, b, c, d), t2: (a, b, d, e), p:(d, N)
output1 = np.zeros((a, b, c, e))
output2 = np.zeros((a, b, c, N))
for i in range(a):
for j in range(b):
output1[i, j, :, :] = np.dot(t1[i, j, :, :], t2[i, j, :, :])
output2[i, j, :, :] = np.dot(t1[i, j, :, :], p)
np.matmulで良い
結論を先にいうとnp.matmulを使うのが最もわかりやすく単純です。$a=6, b=5, c=4, d=3, e=2, N=10$としてみましょう。アフィン変換の前提を無視してt1, t2, pを全て乱数で初期化して計算すると、
t1 = np.random.randn(6, 5, 4, 3)
t2 = np.random.randn(6, 5, 3, 2)
p = np.random.randn(3, 10)
print(np.matmul(t1, t2).shape) # (6, 5, 4, 2)
print(np.matmul(t1, p).shape) # (6, 5, 4, 10)
となり、期待通りの結果になりました。np.matmulを使えばよいのです。
np.dotの場合
np.matmulではなく、np.dotを使った場合どうかというと、テンソルの場合の計算結果が変わります。
t1 = np.random.randn(6, 5, 4, 3)
t2 = np.random.randn(6, 5, 3, 2)
p = np.random.randn(3, 10)
print(np.dot(t1, t2).shape) # (6, 5, 4, 6, 5, 2)
print(np.dot(t1, p).shape) # (6, 5, 4, 10)
matmulの場合はテンソル同士の積は4階テンソルだったのに、dotの場合は6階テンソルになってしまいました。これは、
- (6, 5, 4, 3)の3に注目
- (6, 5, 3, 2)の3に注目
それぞれの積和を取って(6, 5, 4) + (6, 5, 2)→(6, 5, 4, 6, 5, 2)となったわけです。テンソルと行列の場合は同じことをやってもmatmulと同じ結果になります。
np.tensordotの場合
np.tensordotはdotの拡張版です、先程のdotのケースをtensordotで書くと次のようになります。
t1 = np.random.randn(6, 5, 4, 3)
t2 = np.random.randn(6, 5, 3, 2)
p = np.random.randn(3, 10)
print(np.tensordot(t1, t2, [-1, -2]).shape) # (6, 5, 4, 6, 5, 2)
print(np.tensordot(t1, p, [-1, -2]).shape) # (6, 5, 4, 10)
「-1, -2」は消す（和を取る）軸の指定です。tensordotの場合、消す軸を指定することができます。np.dotなら転置しないといけないケース、例えば、
t1 = np.random.randn(6, 5, 4, 3)
t2 = np.random.randn(6, 5, 2, 3)
print(np.dot(t1, t2.swapaxes(-1, -2)).shape) # (6, 5, 4, 6, 5, 2)
print(np.tensordot(t1, t2, [-1, -1]).shape) # (6, 5, 4, 6, 5, 2)
は同じ結果で、tensordotならswapaxesなしでもdotを取れます。swapaxesなしでdotを取るとエラーになります。また、軸が複数のケースもOKで、
t1 = np.random.randn(6, 5, 3, 4)
t2 = np.random.randn(6, 5, 4, 3)
print(np.tensordot(t1, t2, [[-1, -2], [-2, -1]]).shape) # (6, 5, 6, 5)
となります。「具体的にどう使うんだ？」というのはピンとこないですけどね。
まとめ
「アフィン変換をテンソルで計算したい場合はとりあえずnp.matmulを覚えておけばよさそう」ということでした。
Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー