Theano：要素が行列である「行列外積」を取る方法

Question

基本的に、私は2つのテンソルを持つ：A、ここでA.shape = (N, H, D)、B、B.shape = (K, H, D)。私は何をしたいことはテンソルを取得することで、Cは、その形状(N, K, D, H)で：

C[i, j, :, :] = A[i, :, :] * B[j, :, :]. 

これはTheanoで効率的に行うことができますか？

サイドノート：私が達成したい、実際の最終結果は、形状(N, K, D)の、テンソル、Eを持つことがあるように：

E[i, j, :] = (A[i, :, :]*B[j, :, :]).sum(0) 

ので、これを直接取得する方法があれば、私はそれを好むだろう（うまくいけばスペースを節約する）。

Answer 1

一つのアプローチは、使用することを示唆することができbroadcasting -

(A[:,None]*B).sum(2) 

axis=2の和削減前の形状(N, K, H, D)であろう作成された中間配列が(N,K,D).

Answer 2

あなたが得ることができます

にそれを減少させることに注意してください最終的な3次元結果Eは、大きな中間配列を作成せずにbatched_dotを使用します。

import theano.tensor as tt 
A = tt.tensor3('A') # A.shape = (D, N, H) 
B = tt.tensor3('B') # B.shape = (D, H, K) 
E = tt.batched_dot(A, B) # E.shape = (D, N, K) 

残念ながら、これでは入力配列と出力配列の次元を並べ替える必要があります。これはTheanoにdimshuffleで行うことができますが、batched_dotが任意にストライドの配列に対応することができないので、次のことが上げ思わValueError: Some matrix has no unit strideEが評価されている場合：

import theano.tensor as tt 
A = tt.tensor3('A') # A.shape = (N, H, D) 
B = tt.tensor3('B') # B.shape = (K, H, D) 
A_perm = A.dimshuffle((2, 0, 1)) # A_perm.shape = (D, N, H) 
B_perm = B.dimshuffle((2, 1, 0)) # B_perm.shape = (D, H, K) 
E_perm = tt.batched_dot(A_perm, B_perm) # E_perm.shape = (D, N, K) 
E = E_perm.dimshuffle((1, 2, 0)) # E.shape = (N, K, D) 

batched_dotは最初（サイズD）の次元に沿ってscanを使用しています。順次に実行されるので、GPU上で実行されている場合、すべての製品を並列に計算するよりも計算効率が悪くなる可能性があります（scan）。

scanを明示的に使用してブロードキャストのアプローチでbatched_dotアプローチと並列性のメモリ効率をトレードオフすることができます。アイデアは、並行してサイズMのバッチのための完全な製品Cを計算することであろうscanでバッチを反復処理、（Mと仮定するとDの正確な要因である）：

import theano as th 
import theano.tensor as tt 
A = tt.tensor3('A') # A.shape = (N, H, D) 
B = tt.tensor3('B') # B.shape = (K, H, D) 
A_batched = A.reshape((N, H, M, D/M)) 
B_batched = B.reshape((K, H, M, D/M)) 
E_batched, _ = th.scan(
    lambda a, b: (a[:, :, None, :] * b[:, :, :, None]).sum(1), 
    sequences=[A_batched.T, B_batched.T] 
) 
E = E_batched.reshape((D, K, N)).T # E.shape = (N, K, D)