如何在 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

由于dot 规则是“A 的最后一个，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

• 如果任一参数为 N-D，N > 2，则将其视为堆栈 驻留在最后两个索引中的矩阵并相应地广播。

Answer 2

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

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

Related post to understand tensordot.

Answer 3

这是使用 numpy.dot() 的另一种方法，它还可以根据您的要求返回视图，最重要的是比 tensordot 快 4 倍以上 方法，特别是对于小型数组。但是，对于相当大的数组，np.tensordot 比普通的 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() 比 np.tensordot() 快 4 倍

---

# @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() 比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)