如何矢量化此 python 代码以使其更高效？（速度）

Question

我正在尝试使用 numpy 模块对我的代码进行矢量化，原始代码如下：

m = [M[i,:9].dot(N[:9,i]) for i in xrange(9)]

并且我将代码改进为：

m = np.diagonal(M[:9,:9].dot(N[:9,:9]))

但这会导致一些不必要的计算（尤其是当索引远大于9时）。我可以做些什么来进一步提高效率？

编辑：基本上我打算做的是计算两个矩阵 M 和 N 的点积的对角线元素。

Answer 1

您可以使用 np.einsum，因为我们需要保持 M 的第一个轴与 N 的第二个轴对齐，而 reducing/losing 剩余的轴输入。因此，我们将有一个基于 einsum 的解决方案，如下所示：

m = np.einsum('ij,ji->i',M[:,:9],N[:9,:])

Answer 2

假设两个矩阵都是方阵并且具有相同的维度，这就是我们可以做的：

根据矩阵乘法的定义，你要计算的结果基本上是一个点积向量。结果向量的第一个元素是M的第一行与N的第一列的点积，然后是M的第二行与N的第二列的点积，依此类推
我们可以通过转置 N，然后将 M 和 N^T 按元素相乘，然后将每一行相加来表达此计算。这给了我们一个结果列向量。这是代码：
```
Nprime = np.transpose(N[:,:9])
product = np.multiply(M[:9,:], Nprime)
result = np.sum(product, axis=1)
```
根据Divakar的建议，可以浓缩成一行：
```
result = (M[:,:9] * (N[:9,:].T)).sum(1)
```

示例（尺寸为 3 而不是 9）：

import numpy as np

>>> M = np.array(
    [[1, 2, 3],
     [6, 5, 4],
     [8, 7, 9]])

>>> N = np.array(
    [[0, 9, 7],
     [2, 4, 5],
     [6, 8, 1]])

>>> M.dot(N)
array([[22,  41,  20],
       [34, 106,  71],
       [68, 172, 100]])

>>> np.diagonal(M.dot(N))
array([22, 106, 100])
# ^ The reference answer

>>> Nprime = np.transpose(N)
array([[0, 2, 6],
       [9, 4, 8],
       [7, 5, 1]])

>>> product = np.multiply(M, Nprime)
array([[ 0,  4, 18],
       [54, 20, 32],
       [56, 35,  9]])

>>> result = np.sum(product, axis=1)
array([22, 106, 100])

如何矢量化此 python 代码以使其更高效？（速度）

how to vectorize this python code to make it more efficient? (in speed)

python

performance

numpy

vectorization

matrix-multiplication

如何矢量化此 python 代码以使其更高效？ （速度）

how to vectorize this python code to make it more efficient? (in speed)

python

performance

numpy

vectorization

matrix-multiplication

如何矢量化此 python 代码以使其更高效？（速度）