Theano：如何采取元素为矩阵的"matrix outer product"

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)

请注意，正在创建的中间数组在求和减少之前(N, K, H, D)形状axis=2将其减少到 (N,K,D).

Answer 2

您可以使用

batched_dot E获得最终的三维结果：

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无法处理任意步长的数组，因此在评估E时，以下内容会引发ValueError: Some matrix has no unit stride：

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。由于scan是按顺序执行的，如果在 GPU 上运行，这在计算上可能低于并行计算所有产品的效率。

您可以使用显式scan在batched_dot方法的内存效率和广播方法中的并行性之间进行权衡。想法是并行计算大小为M的批次的完整产品C（假设M是D的精确因子），使用scan迭代批处理：

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)