Question

我正在研究基于一些傅里叶变换的模型代码。目前我正在尝试优化它的一部分，以便它可以用于一些大量的数据。在尝试这样做时，我发现了一个奇怪的行为，主要是我的代码的循环版本比使用numpy编写的相同代码更快。测试代码如下：

# -*- coding: utf-8 -*-
import numpy as np
import timeit


def fourier_coef_loop(ts, times, k_max):
    coefs = np.zeros(k_max, dtype=float)
    t = 2.0 * np.pi * (times - times[0]) / times[-1]
    x = np.dot(np.arange(1,k_max+1)[np.newaxis].T,t[np.newaxis])
    for k in xrange(1, k_max + 1):
        cos_k = np.cos(x[k-1])
        coefs[k - 1] = (ts[-1] - ts[0]) + (ts[:-1] * (cos_k[:-1] - cos_k[1:])).sum()
    return(coefs)


def fourier_coef_np(ts, times, k_max):
    coefs = np.zeros(k_max, dtype=float)
    t = 2.0 * np.pi * (times - times[0]) / times[-1]
    x = np.dot(np.arange(1,k_max+1)[np.newaxis].T,t[np.newaxis])
    coefs = np.add(np.einsum('ij,j->i',np.diff(np.cos(x)), -ts[:-1]), (ts[-1] - ts[0]))
    return(coefs)


if __name__ == '__main__':
    iterations = 10
    size = 20000
    setup = "from __main__ import fourier_coef_loop, fourier_coef_np, size\n" \
            "import numpy as np"

#    arg = np.random.normal(size=size)
#    print(np.all(fourier_coef_np(arg, np.arange(size,dtype=float), size / 2) == fourier_coef_loop(arg, np.arange(size,dtype=float), size / 2)))

    time_loop = timeit.timeit("fourier_coef_loop(np.random.normal(size=size), np.arange(size,dtype=float), size / 2)",
                              setup=setup, number=iterations)
    print("With loop: {} s".format(time_loop))

    time_np = timeit.timeit("fourier_coef_np(np.random.normal(size=size), np.arange(size,dtype=float), size / 2)",
                            setup=setup, number=iterations)
    print("With numpy: {} s".format(time_np))

它给出了以下结果：

With loop: 60.8385488987 s
With numpy: 64.9192998409 s

有人可以告诉我为什么循环版本比纯粹的numpy版本更快？我完全没有想法了。我还要感谢有关如何更快地完成此特定功能的任何建议。

Answer 1

正如我在评论中所指出的，我不认为你的时间与循环和矢量化版本之间的差异有关，你的时间甚至包括20000普通伪随机数的生成。如果你想获得精确的图像，你应该尝试在时间之外尽可能多地设置。

无论如何，你的代码有一些奇怪的点，所以这是我自己的建议：

def fourier_coef_new(ts,times,k_max):
    # no need to pre-allocate coefs, you're rebinding later
    t = 2.0 * np.pi * (times - times[0]) / times[-1]
    x = np.arange(1,k_max+1)[:,None] * t # make use of array broadcasting
    coefs = -np.dot(np.diff(np.cos(x)),ts[:-1]) + ts[-1]-ts[0] # einsum was dot
    return(coefs)

我为一组给定的输入测试了这个函数，它给出了与函数相同的结果。请注意[:,None]（或等效[:,np.newaxis]）方式将单例维度引入数组。一旦你有一个形状(N,1)和一个形状(M,)（后者与(1,M)兼容）的数组，根据numpy的数组广播规则，它们的产品将是(N,M)。

我使用给定的，预先生成的数据集快速计时了三个函数，但是在python 3,2中使用size = 2000和3.使用IPython的%timeit构建的在。我并不是说这些结果比你的更可靠，但我怀疑上面的版本应该是最快的：

In [37]: %timeit fourier_coef_loop(ts,times,k_max)
1000 loops, best of 3: 1.09 ms per loop

In [38]: %timeit fourier_coef_np(ts,times,k_max)
1000 loops, best of 3: 1.06 ms per loop

In [39]: %timeit fourier_coef_new(ts,times,k_max)
1000 loops, best of 3: 1.05 ms per loop

正如你所看到的，numpy版本似乎略快一些。由于在两个定时情况下，只有一小部分代码不同（并且两种情况都涉及重三角函数），这似乎是合理的。

python循环比numpy数组操作更快

1 个答案: