python - 沿第三轴的点积

Question

我正在尝试使用 tensordot 在 numpy 中获取张量点积，但我不确定应该如何 reshape 我的数组以实现我的计算。 (总的来说，我对张量的数学还很陌生。)

我有

arr = np.array([[[1, 1, 1],
                [0, 0, 0],
                [2, 2, 2]],

               [[0, 0, 0],
                [4, 4, 4],
                [0, 0, 0]]])

w = [1, 1, 1]

我想沿 axis=2 取点积，这样我就有了矩阵

array([[3, 0, 6],
       [0, 12, 0]])

正确的 numpy 语法是什么？ np.tensordot(arr, [1, 1, 1], axes=2) 似乎引发了一个 ValueError。

Answer 1

arr 沿 axis=2 减少，w 沿 axis=0 减少。因此，使用 np.tensordot ，解决方案是 -

np.tensordot(arr,w,axes=([2],[0]))

或者，也可以使用 np.einsum -

np.einsum('ijk,k->ij',arr,w)

np.matmul也有效

np.matmul(arr, w)

运行时测试-

In [52]: arr = np.random.rand(200,300,300)

In [53]: w = np.random.rand(300)

In [54]: %timeit np.tensordot(arr,w,axes=([2],[0]))
100 loops, best of 3: 8.75 ms per loop

In [55]: %timeit np.einsum('ijk,k->ij',arr,w)
100 loops, best of 3: 9.78 ms per loop

In [56]: %timeit np.matmul(arr, w)
100 loops, best of 3: 9.72 ms per loop

hlin117 在 Macbook Pro OS X El Capitan，numpy 版本 1.10.4 上测试。