【发布时间】:2015-08-21 08:19:38
【问题描述】:
我目前正在使用 GDAL 的 Python 绑定来处理相当大的栅格数据集 (> 4 GB)。由于一次将它们加载到内存中对我来说不是可行的解决方案,我将它们读入较小的块并逐个进行计算。为了避免每次读取的块都分配新的分配,我使用 buf_obj 参数 (here) 将值读入预分配的 NumPy 数组。在某一时刻,我必须计算整个栅格的平均值和标准差。当然,我使用np.std 进行计算。然而,通过分析我的程序的内存消耗,我意识到每次调用 np.std 都会额外分配和释放内存。
演示此行为的最小工作示例:
In [1] import numpy as np
In [2] a = np.random.rand(20e6) # Approx. 150 MiB of memory
In [3] %memit np.mean(a)
peak memory: 187.30 MiB, increment: 0.48 MiB
In [4] %memit np.std(a)
peak memory: 340.24 MiB, increment: 152.91 MiB
在 GitHub 上的 NumPy 源代码树中的搜索显示,np.std 函数在内部调用来自 _methods.py (here) 的 _var 函数。在某一时刻_var 计算与平均值的偏差并将它们相加。因此创建了输入数组的临时副本。该函数本质上计算标准差如下:
mu = sum(arr) / len(arr)
tmp = arr - mu
tmp = tmp * tmp
sd = np.sum(tmp) / len(arr)
虽然这种方法适用于较小的输入数组,但绝对不适用于较大的输入数组。因为我使用的是前面提到的较小的内存块,所以从我的程序的内存角度来看,这个额外的副本并不是一个破坏游戏的问题。然而让我烦恼的是,对于每个块,在读取下一个块之前都会进行新的分配并释放。
在 NumPy 或 SciPy 中是否有其他函数利用诸如 Welford 算法 (Wikipedia) 等具有恒定内存消耗的方法来一次性计算均值和标准差?
另一种方法是实现 _var 函数的自定义版本,并为预分配的缓冲区(如 NumPy ufunc)提供可选的 out 参数。使用这种方法不会消除额外的副本,但至少内存消耗是恒定的,并且每个块中分配的运行时间都被节省了。
编辑:按照 kezzos 的建议测试了 Welford 算法的 Cython 实现。
Cython 实现(从 kezzos 修改):
cimport cython
cimport numpy as np
from libc.math cimport sqrt
@cython.boundscheck(False)
def iterative_approach(np.ndarray[np.float32_t, ndim=1] a):
cdef long n = 0
cdef float mean = 0
cdef float M2 = 0
cdef long i
cdef float delta
cdef float a_min = 10000000 # Must be set to Inf and -Inf for real cases
cdef float a_max = -10000000
for i in range(len(a)):
n += 1
delta = a[i] - mean
mean += delta / n
M2 += delta * (a[i] - mean)
if a[i] < a_min:
a_min = a[i]
if a[i] > a_max:
a_max = a[i]
return a_min, a_max, mean, sqrt(M2 / (n - 1))
NumPy 实现(可以在一个函数中计算均值和标准差):
def vector_approach(a):
return np.min(a), np.max(a), np.mean(a), np.std(a, ddof=1)
使用随机数据集的测试结果(以毫秒为单位,最好是 25 次):
----------------------------------
| Size | Iterative | Vector |
----------------------------------
| 1e2 | 0.00529 | 0.17149 |
| 1e3 | 0.02027 | 0.16856 |
| 1e4 | 0.17850 | 0.23069 |
| 1e5 | 1.93980 | 0.77727 |
| 1e6 | 18.78207 | 8.83245 |
| 1e7 | 180.04069 | 101.14722 |
| 1e8 | 1789.60228 | 1086.66737 |
----------------------------------
似乎使用 Cython 的迭代方法对于较小的数据集更快,而对于具有 10000 多个元素的较大数据集,使用 NumPy 向量(可能是 SIMD 加速)方法似乎更快。所有测试均使用 Python 2.7.9 和 NumPy 1.9.2 版进行。
请注意,在实际情况下,上层函数将用于计算单个栅格块的统计信息。所有区块的标准偏差和平均值将与维基百科 (here) 中建议的方法相结合。它的优点是不需要对栅格的所有元素进行求和,从而避免了浮点溢出问题(至少在某些方面)。
【问题讨论】:
-
您是否尝试过在缓冲区对象上使用 Welford 算法(在纯 Python 中)?
-
@kezzos 我已经用纯 Python 实现的测试结果更新了这个问题。 Python 实现比 NumPy 版本慢很多。