Потребление памяти функции NumPy для стандартного отклонения

В настоящее время я использую привязки Python для GDAL для работы с довольно большими наборами растровых данных ( > 4 ГБ). Поскольку загрузка их в память сразу не представляется возможным, я читаю их на более мелкие блоки и выполняю вычисления по частям. Чтобы избежать нового распределения для каждого чтения блока, я использую аргумент 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

Поиск в дереве источников NumPy на GitHub показал, что функция np.std внутренне вызывает функцию _var из _methods.py (здесь). В какой-то момент _var вычисляются отклонения от среднего и суммируют их. Поэтому создается временная копия входного массива. Функция по существу вычисляет стандартное отклонение следующим образом:

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

Хотя этот подход подходит для небольших массивов ввода, он определенно не подходит для более крупных. Поскольку я использую более мелкие блоки памяти, как упоминалось ранее, эта дополнительная копия не является проблемой для игры с точки зрения памяти в моей программе. Однако то, что меня беспокоит, это то, что для каждого блока новое выделение производится и освобождается перед чтением следующего блока.

Есть ли какая-то другая функция в NumPy или SciPy, которая использует aproach с постоянным потреблением памяти, как алгоритм Welford (Wikipedia) для одного вычисление среднего и стандартного отклонения?

Еще один способ - реализовать пользовательскую версию функции _var с необязательным аргументом out для предварительно распределенного буфера (например, NumPy ufuncs). При таком подходе дополнительная копия не будет устранена, но, по крайней мере, потребление памяти будет постоянным, а время выполнения для распределений в каждом блоке будет сохранено.

РЕДАКТИРОВАТЬ: Протестирована реализация Cython алгоритма Welford, предложенная kezzos.

Реализация 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 (среднее значение и std могут быть вычислены в одной функции):

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, быстрее с меньшими наборами данных, а метод NumPy (возможно, с SIMD-ускорением) подходит для больших наборов данных со 10000+ элементами. Все тесты проводились с Python 2.7.9 и NumPy версии 1.9.2.

Обратите внимание, что в реальном случае верхние функции будут использоваться для вычисления статистики для одного блока растра. Стандартные отклонения и средства для всех блоков должны сочетаться с методологией, предложенной в Википедии (здесь). Преимущество состоит в том, что не все элементы растра необходимо суммировать и тем самым избежать проблемы с переполнением поплавка (по крайней мере, до некоторой точки).

Ответ 1

Китон на помощь! Это достигает хорошей скорости:

%%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 code

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 std

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

Итак, увеличение скорости около 300x. Все еще не так хорошо, как версия numpy.

Ответ 2

Я сомневаюсь, что вы найдете такие функции в numpy. Разница в numpy заключается в том, что он использует команды инструкций vector processor - выполнение одной и той же инструкции больших объемов данных. В основном 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))

Показывая скорость - чем больше blocksize, тем больше скорость. И значительно меньше накладных расходов памяти. Не совсем низкие накладные расходы памяти на 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 [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

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