NumPy函数计算标准差的内存消耗

Question

NumPy函数计算标准差的内存消耗

14

我目前正在使用GDAL的Python绑定来处理相当大的栅格数据集（> 4 GB）。由于一次性将它们加载到内存中不可行，因此我将它们读入较小的块并逐个进行计算。为避免为每个块读取分配新的存储空间，我使用了buf_obj参数(这里)将值读入预先分配的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 中的 _var 函数（这里）。_var 函数在某个时刻计算偏差并将其相加。因此会创建输入数组的临时副本。该函数基本上按照以下方式计算标准偏差：

mu = sum(arr) / len(arr)
tmp = arr - mu
tmp = tmp * tmp
sd = np.sum(tmp) / len(arr)

尽管这种方法对于较小的输入数组可以接受，但对于较大的数组来说绝对不是一个好选择。由于我使用较小的内存块，所以从内存角度来看，这个额外的复制不是我的程序中破坏游戏的问题。然而，让我困扰的是，对于每个块，都会进行新的分配和释放，在读取下一个块之前。

有没有NumPy或SciPy中的其他函数，像Welford算法（维基百科）那样，使用常量内存消耗的方法进行平均值和标准差的单遍计算？

另一种方法是实现_var函数的自定义版本，带有预分配缓冲区的可选out参数（类似于NumPy ufuncs）。采用这种方法，虽然不能消除额外的复制，但至少内存消耗是恒定的，并且节省了每个块中的分配运行时间。

编辑：按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进行。

请注意，在实际情况下，会使用to upper函数来计算栅格的单个块的统计信息。将所有块的标准偏差和均值与维基百科中建议的方法结合起来（此处）。它的优点是不需要对栅格的所有元素进行求和，从而避免了浮点溢出问题（至少在某种程度上）。

- Stefan

你尝试过在缓冲对象上使用Welford算法（纯Python实现）吗？ - kezzos

@kezzos 我已经更新了问题，并提供了一个纯Python实现的测试结果。Python实现比NumPy版本慢得多。 - Stefan

2个回答

5

Cython来拯救！这样可以获得不错的加速效果：

%%cython
cimport cython
cimport numpy as np
from libc.math cimport sqrt

@cython.boundscheck(False)
def std_welford(np.ndarray[np.float64_t, ndim=1] a):
    cdef int n = 0
    cdef float mean = 0
    cdef float M2 = 0
    cdef int a_len = len(a)
    cdef int i
    cdef float delta
    cdef float result
    for i in range(a_len):
        n += 1
        delta = a[i] - mean
        mean += delta / n
        M2 += delta * (a[i] - mean)
    if n < 2:
        result = np.nan
        return result
    else:
        result = sqrt(M2 / (n - 1))
        return result

使用这个来进行测试：

a = np.random.rand(10000).astype(np.float)
print std_welford(a)
%timeit -n 10 -r 10 std_welford(a)

Cython 代码

0.288327455521
10 loops, best of 10: 59.6 µs per loop

原始代码

0.289605617397
10 loops, best of 10: 18.5 ms per loop

Numpy标准差

0.289493223504
10 loops, best of 10: 29.3 µs per loop

因此，速度提高了约300倍。但仍不如numpy版本好。

- kezzos

3

现在你的测试数据和临时变量仍适合存储在CPU缓存中。我认为随着输入数组变得更大，你会发现Numpy表现相对于你的Cython函数越来越差。 - user2379410

当我尝试上面实现的Welford算法时，得到的结果与np.std明显不同。假设arr = np.arange(10, dtype=float)，那么std_welford(arr) == 3.02...，但是np.std(arr) == 2.87...。 - Dunes

1

@Dunes，只是一个定义问题。尝试使用np.std(arr, ddof=1)。 - user2379410

@moarningsun 谢谢，我知道我漏掉了什么。我的统计有点生疏。 - Dunes

1

@kezzos 我已经更新了问题并附上测试结果。它显示数组的大小会影响性能。对于较小的数组，你的实现速度更快，而对于较大的数组，NumPy 更快，尽管所有 NumPy 函数都会多次迭代数组。但是，由于 SIMD 指令集的加速效果似乎超过了此因素。 - Stefan

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Dunes · Accepted Answer

我怀疑你在numpy中找不到这样的函数。 numpy的存在意义在于利用向量处理器指令集--执行大量数据的相同指令。基本上，numpy通过以内存效率换取速度效率。然而，由于Python的内存密集型特性，numpy也能够通过将数据类型与整个数组关联而不是每个单独元素来实现某些内存效率。

提高速度但仍然牺牲一些内存开销的一种方法是分块计算标准差。

import numpy as np

def std(arr, blocksize=1000000):
    """Written for py3, change range to xrange for py2.
    This implementation requires the entire array in memory, but it shows how you can
    calculate the standard deviation in a piecemeal way.
    """
    num_blocks, remainder = divmod(len(arr), blocksize)
    mean = arr.mean()
    tmp = np.empty(blocksize, dtype=float)
    total_squares = 0
    for start in range(0, blocksize*num_blocks, blocksize):
        # get a view of the data we want -- views do not "own" the data they point to
        # -- they have minimal memory overhead
        view = arr[start:start+blocksize]
        # # inplace operations prevent a new array from being created
        np.subtract(view, mean, out=tmp)
        tmp *= tmp
        total_squares += tmp.sum()
    if remainder:
        # len(arr) % blocksize != 0 and need process last part of array
        # create copy of view, with the smallest amount of new memory allocation possible
        # -- one more array *view*
        view = arr[-remainder:]
        tmp = tmp[-remainder:]
        np.subtract(view, mean, out=tmp)
        tmp *= tmp
        total_squares += tmp.sum()
        
    var = total_squares / len(arr)
    sd = var ** 0.5
    return sd

a = np.arange(20e6)
assert np.isclose(np.std(a), std(a))

展示加速度 --- 块大小越大，加速度越大。并且内存开销显著降低。但并非全部降低的内存开销都是100%准确的。

In [70]: %timeit np.std(a)
10 loops, best of 3: 105 ms per loop

In [71]: %timeit std(a, blocksize=4096)
10 loops, best of 3: 160 ms per loop

In [72]: %timeit std(a, blocksize=1000000)
10 loops, best of 3: 105 ms per loop

In [75]: %memit np.std(a)
peak memory: 512.70 MiB, increment: 152.59 MiB

In [73]: %memit std(a, blocksize=4096)
peak memory: 360.11 MiB, increment: 0.00 MiB

In [74]: %memit std(a, blocksize=1000000)
peak memory: 360.11 MiB, increment: 0.00 MiB