Numba - guvectorize比jit快不了多少

11

我试图对许多独立数据集进行蒙特卡洛模拟的并行化处理。我发现numba的并行guvectorize实现仅比numba jit实现快30-40%。

我在Stackoverflow上找到了这些(12)相关主题,但它们并没有真正回答我的问题。在第一个案例中,实现被降级为对象模式而变慢,在第二个案例中,原始发布者未正确使用guvectorize - 这些问题都不适用于我的代码。

为了确保我的代码没有问题,我创建了这个非常简单的代码片段来比较jit和guvectorize:

import timeit
import numpy as np
from numba import jit, guvectorize

#both functions take an (m x n) array as input, compute the row sum, and return the row sums in a (m x 1) array

@guvectorize(["void(float64[:], float64[:])"], "(n) -> ()", target="parallel", nopython=True)
def row_sum_gu(input, output) :
    output[0] = np.sum(input)

@jit(nopython=True)
def row_sum_jit(input_array, output_array) :
    m, n = input_array.shape
    for i in range(m) :
        output_array[i] = np.sum(input_array[i,:])

rows = int(64) #broadcasting (= supposed parallellization) dimension for guvectorize
columns = int(1e6)
input_array = np.ones((rows, columns))
output_array = np.zeros((rows))
output_array2 = np.zeros((rows))

#the first run includes the compile time
row_sum_jit(input_array, output_array)
row_sum_gu(input_array, output_array2)

#run each function 100 times and record the time
print("jit time:", timeit.timeit("row_sum_jit(input_array, output_array)", "from __main__ import row_sum_jit, input_array, output_array", number=100))
print("guvectorize time:", timeit.timeit("row_sum_gu(input_array, output_array2)", "from __main__ import row_sum_gu, input_array, output_array2", number=100))
这给我以下输出(时间有些变化):
jit time: 12.04114792868495
guvectorize time: 5.415564753115177
因此,尽管并行代码利用了所有CPU核心,而jit代码仅使用一个(使用htop验证),但并行代码的性能优势仅为原来的两倍(仅当行数是CPU核心数的整数倍时才会出现,否则性能优势会降低)。
我在一台拥有4个AMD Opteron 6380 CPU(总共64个核心),256GB RAM和Red Hat 4.4.7-1操作系统的机器上运行此代码。我使用Anaconda 4.2.0、Python 3.5.2和Numba 0.26.0。
如何进一步提高并行性能或者我做错了什么?
感谢您的回答。

1
作为参考,我在一台2012年的Macbook Air 1.6 GHz上获得了12秒和3.8秒的时间。因此,尽管你的机器“更好”,但你的JIT时间与我的相同,而你的guvectorize时间更糟糕。 - John Zwinck
你可能还想使用一些随机数据来检查你的两个函数,看看它们是否产生相同的结果。 - JoshAdel
@JoshAdel 我使用了以下代码进行测试:input_array = np.random.rand(rows, columns),并且 np.array_equal(output_array, output_array2) 返回 True。 - Dries Van Laethem
@JohnZwinck 我在不同的机器上运行了代码,有些机器速度较慢,有些机器速度较快,而guvectorize加速在速度较慢的机器上更大,因此我怀疑MSeifert在下面的评论是正确的。 - Dries Van Laethem
@DriesVanLaethem 你是对的。我不确定我最初在测试中做了什么,但现在它们是一致的。我很抱歉。 - JoshAdel
@JoshAdel 没问题,我同意使用随机数而不仅仅是1来测试这种情况是个好主意。 - Dries Van Laethem
1个回答

24

这是因为np.sum太简单了。使用sum处理数组不仅受CPU限制,还受“内存访问”时间的限制。因此,投入更多核心并没有太大区别(当然这取决于与您的CPU相关的内存访问速度有多快)。

只是为了可视化,np.sum大致相当于以下内容(忽略除data之外的任何参数):

def sum(data):
    sum_ = 0.
    data = data.ravel()
    for i in data.size:
        item = data[i]   # memory access (I/O bound)
        sum_ += item     # addition      (CPU bound)
    return sum

因此,如果大部分时间都花在访问内存上,那么并行化不会带来实际的加速效果。但是,如果CPU绑定任务成为瓶颈,那么使用更多的核心将显著加快代码执行速度。

例如,如果您包括一些比加法更慢的操作,您将看到更大的改进:

from math import sqrt
from numba import njit, jit, guvectorize
import timeit
import numpy as np

@njit
def square_sum(arr):
    a = 0.
    for i in range(arr.size):
        a = sqrt(a**2 + arr[i]**2)  # sqrt and square are cpu-intensive!
    return a

@guvectorize(["void(float64[:], float64[:])"], "(n) -> ()", target="parallel", nopython=True)
def row_sum_gu(input, output) :
    output[0] = square_sum(input)

@jit(nopython=True)
def row_sum_jit(input_array, output_array) :
    m, n = input_array.shape
    for i in range(m) :
        output_array[i] = square_sum(input_array[i,:])
    return output_array

我在这里使用了IPython的timeit,但他们应该是等价的:

rows = int(64)
columns = int(1e6)

input_array = np.random.random((rows, columns))
output_array = np.zeros((rows))

# Warmup an check that they are equal 
np.testing.assert_equal(row_sum_jit(input_array, output_array), row_sum_gu(input_array, output_array2))
%timeit row_sum_jit(input_array, output_array.copy())  # 10 loops, best of 3: 130 ms per loop
%timeit row_sum_gu(input_array, output_array.copy())   # 10 loops, best of 3: 35.7 ms per loop

我只使用了4个核心,因此这已经非常接近可能的加速极限了!

请记住,在任务受CPU限制时,并行计算才能显着加速您的计算。


对我来说,仍然令人惊讶的是,OP的“4x AMD Opteron 6380”机器的表现并不比我的2012年款Macbook Air更好,而后者的内存子系统肯定更差(对吧?)。 - John Zwinck
我测试的机器配备了一块Supermicro H8QG6-F主板,带有16个16GB DDR3-1600注册内存。不知道这个配置是比Macbook Air的内存子系统更慢还是更快。 - Dries Van Laethem
@MSeifert 感谢您清晰详细的解释。我正在重构我的数据集和代码,以限制内存访问的次数。 - Dries Van Laethem

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