能否对每次迭代更新变量的numpy for循环实现向量化优化

numpy向量化改造方案 #

完全可以实现向量化改造，你的循环逻辑本质是初始out值加上所有迭代项的累加和，不需要逐次更新复用前序结果，所有迭代项可以批量计算后一次性求和。

改造后代码 #

import numpy as np

x = np.array([[1,2],[2,3],[3,4]])
y = np.array([1,-1,1])
z = np.array([-1,1,-1])
w = 2
out = np.array([10,20])

# 向量化实现
out = out + (w * (y - z)) @ x

验证输出 #

运行后打印out即可得到和原循环完全一致的结果：

>>> out
array([18, 32])

原理解释 #

• 首先计算y - z得到长度为3的一维数组，对应每次迭代的yi - zi
• 乘以标量w后，得到每个迭代的系数w*(yi-zi)，形状为(3,)
• 用矩阵乘法@和形状为(3,2)的x相乘，等价于对x的每一行乘以对应系数后按行求和，直接得到所有迭代项的累加和，形状为(2,)
• 最后加上初始的out数组，就得到最终结果，和循环逻辑完全等价，且全程无Python层循环，性能远高于原始实现。

内容的提问来源于stack exchange，提问作者Mahdi F.