np.zeros 和 np.full 的内存消耗和性能差异的原因

Question

在测量np.zeros的内存消耗时：

import psutil
import numpy as np

process = psutil.Process()
N=10**8
start_rss = process.memory_info().rss
a = np.zeros(N, dtype=np.float64)
print("memory for a", process.memory_info().rss - start_rss)

结果出乎意料 8192 字节，即几乎为 0，而 1e8 双打需要 8e8 字节。

将np.zeros(N, dtype=np.float64) 替换为np.full(N, 0.0, dtype=np.float64) 时，a 所需的内存为800002048 字节。

运行时间也有类似的差异：

import numpy as np
N=10**8
%timeit np.zeros(N, dtype=np.float64)
# 11.8 ms ± 389 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit np.full(N, 0.0, dtype=np.float64)
# 419 ms ± 7.69 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

即np.zeros 对于大尺寸的速度最高可提高 40 倍。

不确定这些差异是否适用于所有架构/操作系统，但我至少在 x86-64 Windows 和 Linux 上观察到了这种差异。

np.zeros和np.full之间的哪些区别可以解释不同的内存消耗和不同的运行时间？

Answer 1

对于这些内存基准测试，我不信任 psutil，而且 rss（驻留集大小）一开始可能不是正确的指标。

使用 stdlib tracemalloc 您可以获得正确的内存分配数字 - 对于这个 N 和 float64 dtype，它应该是大约 800000000 字节的增量：

>>> import numpy as np
>>> import tracemalloc
>>> N = 10**8
>>> tracemalloc.start()
>>> tracemalloc.get_traced_memory()  # current, peak
(159008, 1874350)
>>> a = np.zeros(N, dtype=np.float64)
>>> tracemalloc.get_traced_memory()
(800336637, 802014880)

对于np.full 和np.zeros 之间的时间差异，请比较malloc 和calloc 的手册页，即np.zeros 可以转到an allocation routine which gets zeroed pages。请参阅 PyArray_Zeros --> 调用 PyArray_NewFromDescr_int 传入 1 以获取 zeroed 参数，然后使用 a special case 更快地分配零：

if (zeroed || PyDataType_FLAGCHK(descr, NPY_NEEDS_INIT)) {
    data = npy_alloc_cache_zero(nbytes);
}
else {
    data = npy_alloc_cache(nbytes);
}

看起来np.full 没有这个快速路径。那里的性能将类似于首先执行初始化然后执行复制O(n)：

a = np.empty(N, dtype=np.float64)
a[:] = np.float64(0.0)

numpy 如果填充值为零，开发人员可能已经添加了到np.full 的快速路径，但是为什么还要添加另一种方法来做同样的事情——用户首先可以使用np.zeros。

Answer 2

numpy.zeros 函数直接使用 Numpy 库的 C 代码层，而函数 ones 和 full 通过初始化相同的方式工作一个值数组并在其中复制所需的值。

那么 zeros 函数不需要任何语言解释，而对于其他函数，ones 和 full，Python 代码需要解释为 C 代码。

看看源代码自己弄明白：https://github.com/numpy/numpy/blob/master/numpy/core/numeric.py