Theano：如何在元素是矩阵的情况下采用“矩阵外积”答案

【问题标题】：Theano: How to take a "matrix outer product" where the elements are matricesTheano：如何在元素是矩阵的情况下采用“矩阵外积”
【发布时间】：2016-01-22 16:35:15
【问题描述】：

基本上，我有 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)

所以，如果有办法直接获得它，我会更喜欢它（希望节省空间）。

【问题讨论】：

你想对哪个维度求和？第一个，0?还是“H”，在原始数组中倒数第二个？
在numpy 中可以表示为np.einsum('nhd,khd->nkd', A, B)
我想在 H 上做。假设张量在求和之前的形状为 (1, H, D)，那应该是 sum(1)。
有没有办法在 Theano 中做到这一点？

标签： python numpy matrix theano

【解决方案1】：

可以建议一种使用 broadcasting 的方法 -

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

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

【讨论】：

是的，我最终想到了类似的东西，但我担心它会产生这个巨大的张量，这在 GPU 上工作时可能会令人望而却步。但不确定是否还没有。
@Theo 没错，这里有大量的内存开销。
是的，我想我将从这个解决方案开始，看看它是否适用于我将使用的典型尺寸。谢谢。

【解决方案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 上运行，这可能比并行计算所有产品的计算效率低。

您可以在batched_dot 方法的内存效率和广播方法中显式使用scan 的并行性之间进行权衡。想法是并行计算大小为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)

【讨论】：