在没有执行总和的情况下使用python numpy中的张量积

时间:2014-08-29 19:27:31

标签: python numpy

我有两个张量,每个2D,有一个共同的长轴(例如20.000)和不同的短轴,例如一个9,另一个10)。我想最终得到一个9X10X20000张量,这样对于长轴上的每个位置,其他两个轴都是张量积。 明确地说,这里有“长”轴4,我想做:

A = np.arange(8).reshape(2,4)
B = np.arange(12).reshape(3,4)
C = np.zeros(2,3,4)
for i in range(2):
    for j in range(3):
        for k in range(4):
            C[i,j,k] = A[i,k]*B[j,k]

这段代码有效,但我想知道:有没有一种简单的方法可以做到这一点,而不是为了循环运行?

上下文用于训练神经网络,长轴是训练示例。在计算成本函数的梯度时,我得到了这种形式的公式。 干杯, 利奥

2 个答案:

答案 0 :(得分:3)

我倾向于使用np.einsum来解决这些问题,这样可以很容易地指定在索引方面会发生什么:

>>> A = np.arange(8).reshape(2,4)
>>> B = np.arange(12).reshape(3,4)
>>> np.einsum('ik,jk->ijk', A, B)
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])

答案 1 :(得分:2)

这个特殊的你也可以使用普通的broadcasting

>>> A[:, np.newaxis, :] * B[np.newaxis, :, :]
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])

以上相当于更简单:

>>> A[:, None] * B
array([[[ 0,  1,  4,  9],
        [ 0,  5, 12, 21],
        [ 0,  9, 20, 33]],

       [[ 0,  5, 12, 21],
        [16, 25, 36, 49],
        [32, 45, 60, 77]]])