[英]Improve performance of a for loop in Python (possibly with numpy or numba)
我想在这个函数中改进for
循环的性能。
import numpy as np
import random
def play_game(row, n=1000000):
"""Play the game! This game is a kind of random walk.
Arguments:
row (int[]): row index to use in the p matrix for each step in the
walk. Then length of this array is the same as n.
n (int): number of steps in the random walk
"""
p = np.array([[ 0.499, 0.499, 0.499],
[ 0.099, 0.749, 0.749]])
X0 = 100
Y0 = X0 % 3
X = np.zeros(n)
tempX = X0
Y = Y0
for j in range(n):
tempX = X[j] = tempX + 2 * (random.random() < p.item(row.item(j), Y)) - 1
Y = tempX % 3
return np.r_[X0, X]
困难在于Y
的值是基于X
的值在每个步骤计算的, 并且 Y
然后在下一步骤中用于更新X
的值。
我想知道是否有一些可以产生重大影响的笨拙技巧。 使用Numba是公平的游戏(我试过但没有太大的成功)。 但是,我不想使用Cython。
快速观察告诉我们功能代码中的迭代之间存在数据依赖性。 现在,存在不同种类的数据依赖性。 您正在查看的数据依赖类型是索引依赖性,即在任何迭代中的数据选择取决于先前的迭代计算。 这种依赖似乎很难在迭代之间进行跟踪,因此这篇文章实际上并不是一个矢量化解决方案。 相反,我们会尝试尽可能地预先计算将在循环中使用的值。 基本思想是在循环内做最小的工作。
以下是我们如何进行预先计算的简要说明,从而提供更有效的解决方案:
给定,基于输入row
从中提取行元素的p
的相对较小的形状,可以使用p[row]
从p
预先选择所有那些行。
对于每次迭代,您都在计算一个随机数。 您可以使用可以在循环之前设置的随机数组替换它,因此,您也可以预先计算这些随机值。
根据到目前为止的预先计算的值,您将获得p
所有行的列索引。 请注意,这些列索引将是包含所有可能列索引的大型ndarray,并且在我们的代码中,只有一个将基于每次迭代计算来选择。 使用每次迭代列索引,您可以递增或递减X0
以获得每次迭代输出。
实现看起来像这样 -
randarr = np.random.rand(n)
p = np.array([[ 0.499, 0.419, 0.639],
[ 0.099, 0.749, 0.319]])
def play_game_partvect(row,n,randarr,p):
X0 = 100
Y0 = X0 % 3
signvals = 2*(randarr[:,None] < p[row]) - 1
col_idx = (signvals + np.arange(3)) % 3
Y = Y0
currval = X0
out = np.empty(n+1)
out[0] = X0
for j in range(n):
currval = currval + signvals[j,Y]
out[j+1] = currval
Y = col_idx[j,Y]
return out
要对原始代码进行验证,您可以像这样修改原始代码 -
def play_game(row,n,randarr,p):
X0 = 100
Y0 = X0 % 3
X = np.zeros(n)
tempX = X0
Y = Y0
for j in range(n):
tempX = X[j] = tempX + 2 * (randarr[j] < p.item(row.item(j), Y)) - 1
Y = tempX % 3
return np.r_[X0, X]
请注意,由于此代码预先计算了这些随机值,因此这已经为您提供了比问题中的代码更好的加速。
运行时测试和输出验证 -
In [2]: # Inputs
...: n = 1000
...: row = np.random.randint(0,2,(n))
...: randarr = np.random.rand(n)
...: p = np.array([[ 0.499, 0.419, 0.639],
...: [ 0.099, 0.749, 0.319]])
...:
In [3]: np.allclose(play_game_partvect(row,n,randarr,p),play_game(row,n,randarr,p))
Out[3]: True
In [4]: %timeit play_game(row,n,randarr,p)
100 loops, best of 3: 11.6 ms per loop
In [5]: %timeit play_game_partvect(row,n,randarr,p)
1000 loops, best of 3: 1.51 ms per loop
In [6]: # Inputs
...: n = 10000
...: row = np.random.randint(0,2,(n))
...: randarr = np.random.rand(n)
...: p = np.array([[ 0.499, 0.419, 0.639],
...: [ 0.099, 0.749, 0.319]])
...:
In [7]: np.allclose(play_game_partvect(row,n,randarr,p),play_game(row,n,randarr,p))
Out[7]: True
In [8]: %timeit play_game(row,n,randarr,p)
10 loops, best of 3: 116 ms per loop
In [9]: %timeit play_game_partvect(row,n,randarr,p)
100 loops, best of 3: 14.8 ms per loop
因此,我们看到加速大约7.5x+
,不错!
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.