python - 自定义numpy的向量化运算的效率问题

Question

我有一个 python 函数如下:

def myfun(x):
    if x > 0:
        return 0
    else:
        return np.exp(x)

哪里np是 numpy图书馆。我想在 numpy 中使函数矢量化，所以我使用:

vec_myfun = np.vectorize(myfun)

我做了一个测试来评估效率。首先，我生成一个包含 100 个随机数的向量:

x = np.random.randn(100)

然后我运行以下代码来获取运行时:

%timeit np.exp(x)

%timeit vec_myfun(x)

np.exp(x) 的运行时间是1.07 µs ± 24.9 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each) .

vec_myfun(x) 的运行时间是71.2 µs ± 1.68 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

我的问题是:与 np.exp 相比, vec_myfun只有一个额外的步骤来检查 $x$ 的值，但它比 np.exp 运行得慢得多.有没有一种有效的矢量化方式 myfun使其与np.exp一样高效？

Answer 1

使用np.where:

>>> x = np.random.rand(100,)
>>> %timeit np.exp(x)
1.22 µs ± 49.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
>>> %timeit np.where(x > 0, 0, np.exp(x))
4.09 µs ± 282 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

相比之下，您的向量化函数在我的机器上运行时间约为 30 微秒。

至于为什么运行的慢，只是比np.exp复杂多了。它进行了大量的类型推导、广播，并且可能对实际方法进行多次调用。其中大部分发生在 Python 本身，而调用 np.exp(以及此处的 np.where 版本)的几乎所有内容都在 C 中。