Numba 和 numpy 数组分配：为什么这么慢？
Numba and numpy array allocation: why is it so slow?
我最近使用 Cython 和 Numba 来加速进行数值模拟的 python 的小部分。起初，使用 numba 进行开发似乎更容易。然而，我发现很难理解 numba 何时会提供更好的性能，何时不会。
性能意外下降的一个例子是当我使用函数 np.zeros() 在编译函数中分配一个大数组时。例如，考虑三个函数定义：
import numpy as np
from numba import jit
def pure_python(n):
mat = np.zeros((n,n), dtype=np.double)
# do something
return mat.reshape((n**2))
@jit(nopython=True)
def pure_numba(n):
mat = np.zeros((n,n), dtype=np.double)
# do something
return mat.reshape((n**2))
def mixed_numba1(n):
return mixed_numba2(np.zeros((n,n)))
@jit(nopython=True)
def mixed_numba2(array):
n = len(array)
# do something
return array.reshape((n,n))
# To compile
pure_numba(10)
mixed_numba1(10)
由于 #do something 是空的，我不希望 pure_numba 函数更快。然而，我没想到会出现这样的性能下降：
n=10000
%timeit x = pure_python(n)
%timeit x = pure_numba(n)
%timeit x = mixed_numba1(n)
我在 mac 上获得 (python 3.7.7, numba 0.48.0)
4.96 µs ± 65.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
344 ms ± 7.76 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
3.8 µs ± 30.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
在这里，当我在编译函数中使用函数 np.zeros() 时，numba 代码要慢得多。 np.zeros() 在函数外时正常。
我是不是做错了什么，或者我应该总是像这些由 numba 编译的外部函数那样分配大数组？
更新
当 n 足够大时，这似乎与 np.zeros((n,n)) 对矩阵的惰性初始化有关（参见 ）。
for n in [1000, 2000, 5000]:
print('n=',n)
%timeit x = pure_python(n)
%timeit x = pure_numba(n)
%timeit x = mixed_numba1(n)
给我：
n = 1000
468 µs ± 15.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
296 µs ± 6.55 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
300 µs ± 2.26 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
n = 2000
4.79 ms ± 182 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.45 ms ± 36 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.54 ms ± 127 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
n = 5000
270 µs ± 4.66 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
104 ms ± 599 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
119 µs ± 1.24 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
tl;dr Numpy 使用 C 内存函数，而 Numba 必须分配零
我写了一个脚本来绘制完成多个选项所需的时间，当 np.zeros 数组的大小达到我的 2048*2048*8 = 32 MB 时，Numba 的性能似乎严重下降机器如下图所示。
Numba 的 np.zeros 实现与创建一个空数组并通过迭代数组的维度用零填充它一样快（这是 Numba 嵌套循环 图中的绿色曲线）。这实际上可以通过在 运行 脚本之前设置 NUMBA_DUMP_IR 环境变量来进行双重检查（见下文）。与 numba_loop 的转储相比，差别不大。
有趣的是，np.zeros 在超过 32 MB 阈值时略有提升。
尽管我远非专家，但我的最佳猜测是 32 MB 限制是一个 OS 或硬件瓶颈，来自同一进程缓存中可容纳的数据量。如果超过这个，将数据移入和移出缓存以对其进行操作的操作非常耗时。
相比之下，Numpy 使用 calloc 获取一些内存段，并承诺在访问数据时用零填充数据。
这就是我的进展，我意识到这只是答案的一半，但也许更有知识的人可以阐明实际发生的事情。
Numba IR 转储：
---------------------------IR DUMP: pure_numba_zeros----------------------------
label 0:
n = arg(0, name=n) ['n']
load_global.0 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['load_global.0']
load_attr.1 = getattr(value=load_global.0, attr=zeros) ['load_global.0', 'load_attr.1']
del load_global.0 []
build_tuple.4 = build_tuple(items=[Var(n, script.py:15), Var(n, script.py:15)]) ['build_tuple.4', 'n', 'n']
load_global.5 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['load_global.5']
load_attr.6 = getattr(value=load_global.5, attr=double) ['load_global.5', 'load_attr.6']
del load_global.5 []
call_function_kw.8 = call load_attr.1(build_tuple.4, func=load_attr.1, args=[Var(build_tuple.4, script.py:15)], kws=[('dtype', Var(load_attr.6, script.py:15))], vararg=None) ['build_tuple.4', 'load_attr.6', 'call_function_kw.8', 'load_attr.1']
del load_attr.1 []
del load_attr.6 []
del build_tuple.4 []
mat = call_function_kw.8 ['call_function_kw.8', 'mat']
del call_function_kw.8 []
load_method.10 = getattr(value=mat, attr=reshape) ['load_method.10', 'mat']
del mat []
$const28.12 = const(int, 2) ['$const28.12']
binary_power.13 = n ** $const28.12 ['binary_power.13', '$const28.12', 'n']
del n []
del $const28.12 []
call_method.14 = call load_method.10(binary_power.13, func=load_method.10, args=[Var(binary_power.13, script.py:16)], kws=(), vararg=None) ['load_method.10', 'binary_power.13', 'call_method.14']
del binary_power.13 []
del load_method.10 []
return_value.15 = cast(value=call_method.14) ['call_method.14', 'return_value.15']
del call_method.14 []
return return_value.15 ['return_value.15']
生成图表的脚本：
import numpy as np
from numba import jit
from time import time
import os
import matplotlib.pyplot as plt
os.environ['NUMBA_DUMP_IR'] = '1'
def numpy_zeros(n):
mat = np.zeros((n,n), dtype=np.double)
return mat.reshape((n**2))
@jit(nopython=True)
def numba_zeros(n):
mat = np.zeros((n,n), dtype=np.double)
return mat.reshape((n**2))
@jit(nopython=True)
def numba_loop(n):
mat = np.empty((n * 2,n), dtype=np.float32)
for i in range(mat.shape[0]):
for j in range(mat.shape[1]):
mat[i, j] = 0.
return mat.reshape((2 * n**2))
# To compile
numba_zeros(10)
numba_loop(10)
os.environ['NUMBA_DUMP_IR'] = '0'
max_n = 4100
time_deltas = {
'numpy_zeros': [],
'numba_zeros': [],
'numba_loop': [],
}
call_count = 10
for n in range(0, max_n, 10):
for f in (numpy_zeros, numba_zeros, numba_loop):
start = time()
for i in range(call_count):
x = f(n)
delta = time() - start
time_deltas[f.__name__].append(delta / call_count)
print(f'{f.__name__:25} n = {n}: {delta}')
print()
size = np.arange(0, max_n, 10) ** 2 * 8 / 1024 ** 2
fig, ax = plt.subplots()
plt.xticks(np.arange(0, size[-1], 16))
plt.axvline(x=32, color='gray', lw=0.5)
ax.plot(size, time_deltas['numpy_zeros'], label='Numpy zeros (calloc)')
ax.plot(size, time_deltas['numba_zeros'], label='Numba zeros')
ax.plot(size, time_deltas['numba_loop'], label='Numba nested loop')
ax.set_xlabel('Size of array in MB')
ax.set_ylabel(r'Mean $\Delta$t in s')
plt.legend(loc='upper left')
plt.show()
我最近使用 Cython 和 Numba 来加速进行数值模拟的 python 的小部分。起初，使用 numba 进行开发似乎更容易。然而，我发现很难理解 numba 何时会提供更好的性能，何时不会。
性能意外下降的一个例子是当我使用函数 np.zeros() 在编译函数中分配一个大数组时。例如，考虑三个函数定义：
import numpy as np
from numba import jit
def pure_python(n):
mat = np.zeros((n,n), dtype=np.double)
# do something
return mat.reshape((n**2))
@jit(nopython=True)
def pure_numba(n):
mat = np.zeros((n,n), dtype=np.double)
# do something
return mat.reshape((n**2))
def mixed_numba1(n):
return mixed_numba2(np.zeros((n,n)))
@jit(nopython=True)
def mixed_numba2(array):
n = len(array)
# do something
return array.reshape((n,n))
# To compile
pure_numba(10)
mixed_numba1(10)
由于 #do something 是空的，我不希望 pure_numba 函数更快。然而，我没想到会出现这样的性能下降：
n=10000
%timeit x = pure_python(n)
%timeit x = pure_numba(n)
%timeit x = mixed_numba1(n)
我在 mac 上获得 (python 3.7.7, numba 0.48.0)
4.96 µs ± 65.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
344 ms ± 7.76 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
3.8 µs ± 30.3 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
在这里，当我在编译函数中使用函数 np.zeros() 时，numba 代码要慢得多。 np.zeros() 在函数外时正常。
我是不是做错了什么，或者我应该总是像这些由 numba 编译的外部函数那样分配大数组？
更新
当 n 足够大时，这似乎与 np.zeros((n,n)) 对矩阵的惰性初始化有关（参见
for n in [1000, 2000, 5000]:
print('n=',n)
%timeit x = pure_python(n)
%timeit x = pure_numba(n)
%timeit x = mixed_numba1(n)
给我：
n = 1000
468 µs ± 15.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
296 µs ± 6.55 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
300 µs ± 2.26 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
n = 2000
4.79 ms ± 182 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.45 ms ± 36 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.54 ms ± 127 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
n = 5000
270 µs ± 4.66 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
104 ms ± 599 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
119 µs ± 1.24 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
tl;dr Numpy 使用 C 内存函数，而 Numba 必须分配零
我写了一个脚本来绘制完成多个选项所需的时间，当 np.zeros 数组的大小达到我的 2048*2048*8 = 32 MB 时，Numba 的性能似乎严重下降机器如下图所示。
Numba 的 np.zeros 实现与创建一个空数组并通过迭代数组的维度用零填充它一样快（这是 Numba 嵌套循环 图中的绿色曲线）。这实际上可以通过在 运行 脚本之前设置 NUMBA_DUMP_IR 环境变量来进行双重检查（见下文）。与 numba_loop 的转储相比，差别不大。
有趣的是，np.zeros 在超过 32 MB 阈值时略有提升。
尽管我远非专家，但我的最佳猜测是 32 MB 限制是一个 OS 或硬件瓶颈，来自同一进程缓存中可容纳的数据量。如果超过这个，将数据移入和移出缓存以对其进行操作的操作非常耗时。
相比之下，Numpy 使用 calloc 获取一些内存段，并承诺在访问数据时用零填充数据。
这就是我的进展，我意识到这只是答案的一半，但也许更有知识的人可以阐明实际发生的事情。
Numba IR 转储：
---------------------------IR DUMP: pure_numba_zeros----------------------------
label 0:
n = arg(0, name=n) ['n']
load_global.0 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['load_global.0']
load_attr.1 = getattr(value=load_global.0, attr=zeros) ['load_global.0', 'load_attr.1']
del load_global.0 []
build_tuple.4 = build_tuple(items=[Var(n, script.py:15), Var(n, script.py:15)]) ['build_tuple.4', 'n', 'n']
load_global.5 = global(np: <module 'numpy' from '/lib/python3.8/site-packages/numpy/__init__.py'>) ['load_global.5']
load_attr.6 = getattr(value=load_global.5, attr=double) ['load_global.5', 'load_attr.6']
del load_global.5 []
call_function_kw.8 = call load_attr.1(build_tuple.4, func=load_attr.1, args=[Var(build_tuple.4, script.py:15)], kws=[('dtype', Var(load_attr.6, script.py:15))], vararg=None) ['build_tuple.4', 'load_attr.6', 'call_function_kw.8', 'load_attr.1']
del load_attr.1 []
del load_attr.6 []
del build_tuple.4 []
mat = call_function_kw.8 ['call_function_kw.8', 'mat']
del call_function_kw.8 []
load_method.10 = getattr(value=mat, attr=reshape) ['load_method.10', 'mat']
del mat []
$const28.12 = const(int, 2) ['$const28.12']
binary_power.13 = n ** $const28.12 ['binary_power.13', '$const28.12', 'n']
del n []
del $const28.12 []
call_method.14 = call load_method.10(binary_power.13, func=load_method.10, args=[Var(binary_power.13, script.py:16)], kws=(), vararg=None) ['load_method.10', 'binary_power.13', 'call_method.14']
del binary_power.13 []
del load_method.10 []
return_value.15 = cast(value=call_method.14) ['call_method.14', 'return_value.15']
del call_method.14 []
return return_value.15 ['return_value.15']
生成图表的脚本：
import numpy as np
from numba import jit
from time import time
import os
import matplotlib.pyplot as plt
os.environ['NUMBA_DUMP_IR'] = '1'
def numpy_zeros(n):
mat = np.zeros((n,n), dtype=np.double)
return mat.reshape((n**2))
@jit(nopython=True)
def numba_zeros(n):
mat = np.zeros((n,n), dtype=np.double)
return mat.reshape((n**2))
@jit(nopython=True)
def numba_loop(n):
mat = np.empty((n * 2,n), dtype=np.float32)
for i in range(mat.shape[0]):
for j in range(mat.shape[1]):
mat[i, j] = 0.
return mat.reshape((2 * n**2))
# To compile
numba_zeros(10)
numba_loop(10)
os.environ['NUMBA_DUMP_IR'] = '0'
max_n = 4100
time_deltas = {
'numpy_zeros': [],
'numba_zeros': [],
'numba_loop': [],
}
call_count = 10
for n in range(0, max_n, 10):
for f in (numpy_zeros, numba_zeros, numba_loop):
start = time()
for i in range(call_count):
x = f(n)
delta = time() - start
time_deltas[f.__name__].append(delta / call_count)
print(f'{f.__name__:25} n = {n}: {delta}')
print()
size = np.arange(0, max_n, 10) ** 2 * 8 / 1024 ** 2
fig, ax = plt.subplots()
plt.xticks(np.arange(0, size[-1], 16))
plt.axvline(x=32, color='gray', lw=0.5)
ax.plot(size, time_deltas['numpy_zeros'], label='Numpy zeros (calloc)')
ax.plot(size, time_deltas['numba_zeros'], label='Numba zeros')
ax.plot(size, time_deltas['numba_loop'], label='Numba nested loop')
ax.set_xlabel('Size of array in MB')
ax.set_ylabel(r'Mean $\Delta$t in s')
plt.legend(loc='upper left')
plt.show()