如何对这个 python 代码进行矢量化以提高效率？（速度）答案

【问题标题】：how to vectorize this python code to make it more efficient? (in speed)如何对这个 python 代码进行矢量化以提高效率？（速度）
【发布时间】：2016-11-20 12:34:25
【问题描述】：

我正在尝试使用 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的点积的对角线元素。

【问题讨论】：

提及代码实际想要完成的工作总是很高兴。
感谢您的建议，我现在重新编辑了我的问题
对不起，误导信息，但在我的实际代码中，两个矩阵的形状不一定是正方形。

标签： python performance numpy vectorization matrix-multiplication

【解决方案1】：

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

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

【讨论】：

感谢您的建议，确实提高了数值效率！

【解决方案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])

【讨论】：

大多数情况下矩阵不是方阵，但感谢您的建议！
@M.Gu 我认为这仍然适用于非方阵。