如何在 numpy 中沿特定维度取点积？

Question

我有两个数组。一个是 n by p，另一个是 d by p by r。我希望我的输出是 d by n by r，我可以在下面构造张量 B 时轻松实现。但是，我想在没有该循环的情况下执行此操作。

import numpy

X = numpy.array([[1,2,3],[3,4,5],[5,6,7],[7,8,9]]) # n x p
betas = numpy.array([[[1,2],[1,2],[1,2]], [[5,6],[5,6],[5,6]]]) # d x p x r

print X.shape
print betas.shape

B = numpy.zeros((betas.shape[0],X.shape[0],betas.shape[2]))
print B.shape

for i in range(B.shape[0]):
    B[i,:,:] = numpy.dot(X, betas[i])

print "B",B

C = numpy.tensordot(X, betas, axes=([1],[0]))
print C.shape

我尝试过各种方法让C匹配B，但到目前为止我一直没有成功。有没有不涉及调用 reshape 的方法？

Answer 1

我们可以使用 np.tensordot 然后需要置换轴 -

B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)
# Or np.tensordot(X, betas, axes=(1,1)).swapaxes(0,1)

.

Answer 2

由于 dot 规则是 'last of A with 2nd to the last of B'，您可以执行 X.dot(betas) 并得到一个 (n,d,r) 数组（这对共享的 p尺寸）。然后你只需要一个转置就可以得到 (d,n,r)

In [200]: X.dot(betas).transpose(1,0,2)
Out[200]: 
array([[[  6,  12],
        [ 12,  24],
        [ 18,  36],
        [ 24,  48]],

       [[ 30,  36],
        [ 60,  72],
        [ 90, 108],
        [120, 144]]])

我们也可以直接根据尺寸规范写出einsum版本：

np.einsum('np,dpr->dnr', X,betas)

matmul 也是如此（这在最后 2 个轴上是 dot，而 d 则随行）。

X@betas

If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes and broadcast accordingly.

Answer 3

这是另一种使用 numpy.dot(), which also returns a view as you requested, and most importantly more than 4x faster than tensordot approach, particularly for small sized arrays. But, np.tensordot is way faster than plain np.dot() 的方法，用于相当大的数组。请参阅下面的时间安排。

In [108]: X.shape
Out[108]: (4, 3)

In [109]: betas.shape
Out[109]: (2, 3, 2)

# use `np.dot` and roll the second axis to first position
In [110]: dot_prod = np.rollaxis(np.dot(X, betas), 1)

In [111]: dot_prod.shape
Out[111]: (2, 4, 2)

# @Divakar's approach
In [113]: B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)

# sanity check :)
In [115]: np.all(np.equal(dot_prod, B))
Out[115]: True

现在，两种方法的性能：

对于小型数组 np.dot() is 4x faster than np.tensordot()

# @Divakar's approach
In [117]: %timeit B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)
10.6 µs ± 2.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# @hpaulj's approach
In [151]: %timeit esum_dot = np.einsum('np, dpr -> dnr', X, betas)
4.16 µs ± 235 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

# proposed approach: more than 4x faster!!
In [118]: %timeit dot_prod = np.rollaxis(np.dot(X, betas), 1)
2.47 µs ± 11.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

对于相当大的阵列，np.tensordot() is much faster than np.dot()

In [129]: X = np.random.randint(1, 10, (600, 500))
In [130]: betas = np.random.randint(1, 7, (300, 500, 300))

In [131]: %timeit B = np.tensordot(betas, X, axes=(1,1)).swapaxes(1,2)
18.2 s ± 2.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [132]: %timeit dot_prod = np.rollaxis(np.dot(X, betas), 1)
52.8 s ± 14.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

如何在 numpy 中沿特定维度取点积？

How can I take a dot product along a particular dimension in numpy?

python

numpy

linear-algebra

multidimensional-array

numpy-ndarray