使用元素指数加速嵌套for循环

6

我正在处理一段庞大的代码,发现需要加速其中的某个部分。我已经创建了下面展示的 MWE

import numpy as np
import time

def random_data(N):
    # Generate some random data.
    return np.random.uniform(0., 10., N).tolist()

# Lists that contain all the data.
list1 = [random_data(10) for _ in range(1000)]
list2 = [random_data(1000), random_data(1000)]

# Start taking the time.
tik = time.time()

list4 = []
# Loop through all elements in list1.
for elem in list1:

    list3 = []
    # Loop through elements in list2.
    for elem2 in zip(*list2):

        A = np.exp(-0.5*((elem[0]-elem2[0])/elem[3])**2)
        B = np.exp(-0.5*((elem[1]-elem2[1])/elem[3])**2)
        list3.append(A*B)

    # Sum elements in list3 and append result to list4.
    sum_list3 = sum(list3) if sum(list3)>0. else 1e-06
    list4.append(sum_list3)

# Print the elapsed time.
print time.time()-tik
< p > list1list2 的奇怪格式是因为这个代码块需要它们。

显而易见的部分,大部分时间都花在了递归计算AB项上。

有没有办法加速这段代码的执行而不必进行并行化(我以前尝试过,但给我带来了很多 麻烦)?我可以使用任何包,例如numpyscipy等。


添加

这是应用abarnert的优化和Jaime的建议后的结果,只进行了一次指数运算。在我的系统上,优化后的函数平均快了约60倍。

import numpy as np
import timeit

def random_data(N):
    return np.random.uniform(0., 10., N).tolist()

# Lists that contain all the data.
list1 = [random_data(10) for _ in range(1000)]
list2 = [random_data(1000), random_data(1000)]

array1 = np.array(list1)
array2 = np.array(zip(*list2))


# Old non-optimezed function.
def func1():
    list4 = []
    # Process all elements in list1.
    for elem in list1:
        # Process all elements in list2.
        list3 = []
        for elem2 in zip(*list2):
            A = np.exp(-0.5*((elem[0]-elem2[0])/elem[3])**2)
            B = np.exp(-0.5*((elem[1]-elem2[1])/elem[3])**2)
            list3.append(A*B)
        # Sum elements in list3 and append result to list4.
        sum_list3 = sum(list3) if sum(list3)>0. else 1e-06
        list4.append(sum_list3)

# New optimized function.
def func2():
    list4 = []
    # Process all elements in list1.
    for elem in array1:

        # Broadcast over elements in array2.
        A = -0.5*((elem[0]-array2[:,0])/elem[3])**2
        B = -0.5*((elem[1]-array2[:,1])/elem[3])**2
        array3 = np.exp(A+B)

        # Sum elements in array3 and append result to list4.
        sum_list3 = max(array3.sum(), 1e-10)
        list4.append(sum_list3)


# Get time for both functions.
func1_time = timeit.timeit(func1, number=10)
func2_time = timeit.timeit(func2, number=10)

# Print hom many times faster func2 is versus func1.
print func1_time/func2_time

2
为什么这段代码在顶部有NumPy依赖项,但却全部基于列表? - user2357112
没有特别的原因,这段代码只是从另一段代码中接收了list1list2。至于list3list4,这是我能想到的最好的填充方法。如果您认为有所区别,它们都可以转换为numpy数组。 - Gabriel
2
@Gabriel:当然会有所不同。这就是使用numpy的全部意义——如果您可以在数组上广播计算,那么您将用C循环替换Python循环,并消除每个算术计算周围的所有装箱/拆箱,这意味着您的代码通常会快4-400倍。 - abarnert
3
发一个好的 MWE,点赞!这是提问并得到了极好回答的努力工作的绝佳示范。我会收藏它以便将来链接。 - YXD
1
最后一个附带说明:不要尝试使用 time.time 自己计时。 [timeit] (http://docs.python.org/3/library/timeit.html) 模块(如果您在使用IPython,则使用%timeit魔术语句)会确保选择正确的计时器,处理您甚至没有想到的许多问题,允许您重复测试并正确汇总它们,并使事情更简单。 (当您的代码比预期慢100倍时,通常不是太大的问题,但值得养成始终使用 timeit 的习惯。) - abarnert
你说得完全正确,到目前为止我从未使用过那个模块,但现在我会尝试应用它来计算优化后的函数快了多少倍。我会将这个问题补充上去。 - Gabriel
1个回答

7
你希望逐步将这个过程从使用列表和循环转换为使用数组和广播,首先抓住最容易或最关键的部分,直到它变得足够快。
第一步是不要反复执行那个`zip(*list2)`(特别是在Python 2.x中)。顺便说一下,我们可以将其存储在一个数组中,并对`list1`执行相同的操作 - 目前仍然可以迭代它们。所以:
array1 = np.array(list1)
array2 = np.array(zip(*list2))
# …
for elem in array1:
    # …
    for elem2 in array2:

这样做并不能显著提高速度,在我的机器上,它只能将运行时间从14.1秒缩短到12.9秒,但是它为我们提供了一个开始工作的地方。

你还应该删除对sum(list3)的重复计算:

sum_list3 = sum(list3)
sum_list3 = sum_list3 if sum_list3>0. else 1e-06

同时,你希望value <= 0变成1e-6,但0 < value < 1e-6不变有点奇怪。这是故意的吗?如果不是,你可以通过以下方式修复它,并简化代码:

sum_list3 = max(array3.sum(), 1e-06)

现在,让我们广播AB的计算结果:
# Broadcast over elements in list2.
A = np.exp(-0.5*((elem[0]-array2[:,0])/elem[3])**2)
B = np.exp(-0.5*((elem[1]-array2[:, 1])/elem[3])**2)
array3 = A*B

# Sum elements in list3 and append result to list4.
sum_list3 = max(array3.sum(), 1e-06)

list4.append(sum_list3)

这将我们的时间从12.9秒降到了0.12秒。您还可以进一步通过在array1上进行广播,并用预先分配的数组替换list4等来提高性能,但这已经足够快了。


4
指数运算代价高昂:不要使用两次 np.exp 然后相乘,而是将这两个值相加然后再使用 np.exp - Jaime
1
@Jaime:我们已经获得了99%的加速,我怀疑他不需要再提高10%。(我快速估算了一下-跳过exp可以将时间缩短到0.79倍,因此只做一个而不是两个可能约为0.9倍。)但我敢打赌,通过将其移出一层(广播exp到所有行,然后对所有行进行sum,而不是每行都进行一次),你可以获得更多的改进。因此,如果现有的改进还不够,首先应该在array1array4上进行广播,然后再优化它们内部的数学运算。 - abarnert
2
如果你将整个向量化,那么可以获得额外的20%,我必须承认,在前99%上这是微不足道的。但这是我最讨厌的事情之一,就像如果平方已经足够了,不需要再开方距离一样,我忍不住要说一下,抱歉给大家带来干扰... - Jaime
1
@Jaime:这不是噪音,将额外的优化放在评论中是值得的——以防万一原帖需要它们,同时也是为了让原帖从中学习。我只是解释了为什么我认为它不需要成为答案的一部分,作为评论可能已经足够好了。 - abarnert
2
@Gabriel:人们经常不加思考地说这句话,但请考虑一下:为了每次运行节省 180 微秒的时间,值得花费 2 小时来加速你的代码吗?如果这样做了 10000 次,你最多只能节省负 3598.2 秒——而且你增加了出错的可能性,以及调试所需的时间(特别是如果你发现新版本更难理解)。 - abarnert
显示剩余6条评论

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接