Использование памяти numpy memmap - хотите повторить один раз

Question

Скажем, у меня есть большая матрица, сохраненная на диске. хранить все это в памяти не реально, поэтому я использую memmap для доступа к нему

A = np.memmap(filename, dtype='float32', mode='r', shape=(3000000,162))

Теперь допустим, что я хочу перебрать эту матрицу (не по существу в упорядоченном виде) так, чтобы каждая строка была доступна ровно один раз.

p = some_permutation_of_0_to_2999999()

Я хотел бы сделать что-то подобное:

start = 0
end = 3000000
num_rows_to_load_at_once = some_size_that_will_fit_in_memory()
while start < end:
    indices_to_access = p[start:start+num_rows_to_load_at_once]
    do_stuff_with(A[indices_to_access, :])
    start = min(end, start+num_rows_to_load_at_once)

по мере того, как этот процесс идет на моем компьютере, становится все медленнее и медленнее, а объем используемой оперативной и виртуальной памяти растет

Есть ли способ заставить np.memmap использовать до определенного объема памяти? (Я знаю, что мне не понадобится больше, чем количество строк, которые я планирую прочитать за раз, и это кэширование не поможет мне, так как я получаю доступ к каждой строке ровно один раз)

Может быть, вместо этого есть какой-то другой способ итерации (как генератор) по массиву np в произвольном порядке? Я мог бы написать это вручную, используя file.seek, но это происходит намного медленнее, чем реализация np.memmap.

do_stuff_with() не сохраняет никаких ссылок на массив, который он получает, поэтому никаких "утечек памяти" в этом аспекте

Спасибо

Answer 1

Это была проблема, с которой я некоторое время пытался разобраться. Я работаю с большими наборами данных изображений, и numpy.memmap предлагает удобное решение для работы с этими большими наборами.

Однако, как вы указали, если мне понадобится получить доступ к каждому кадру (или строке в вашем случае) для выполнения какой-либо операции, использование ОЗУ в конечном итоге будет максимальным.

К счастью, я недавно нашел решение, которое позволит вам перебирать весь массив memmap, ограничивая использование оперативной памяти.

Решение:

import numpy as np

# create a memmap array
input = np.memmap('input', dtype='uint16', shape=(10000,800,800), mode='w+')

# create a memmap array to store the output
output = np.memmap('output', dtype='uint16', shape=(10000,800,800), mode='w+')

def iterate_efficiently(input, output, chunk_size):
    # create an empty array to hold each chunk
    # the size of this array will determine the amount of RAM usage
    holder = np.zeros([chunk_size,800,800], dtype='uint16')

    # iterate through the input, replace with ones, and write to output
    for i in range(input.shape[0]):
        if i % chunk_size == 0:
            holder[:] = input[i:i+chunk_size] # read in chunk from input
            holder += 5 # perform some operation
            output[i:i+chunk_size] = holder # write chunk to output

def iterate_inefficiently(input, output):
    output[:] = input[:] + 5

Сроки Результаты:

In [11]: %timeit iterate_efficiently(input,output,1000)
1 loop, best of 3: 1min 48s per loop

In [12]: %timeit iterate_inefficiently(input,output)
1 loop, best of 3: 2min 22s per loop

Размер массива на диске составляет ~12 ГБ. С использованием iterate_efficiently Функция сохраняет использование памяти до 1,28 ГБ, тогда как iterate_inefficiently функция в конечном итоге достигает 12 ГБ в оперативной памяти.

Это было проверено на Mac OS.

Answer 2

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

Вариант 1 - реконструировать карту памяти при каждом чтении / записи:

def MoveMMapNPArray(data, output_filename):
    CHUNK_SIZE = 4096
    for idx in range(0,x.shape[1],CHUNK_SIZE):
        x = np.memmap(data.filename, dtype=data.dtype, mode='r', shape=data.shape, order='F')
        y = np.memmap(output_filename, dtype=data.dtype, mode='r+', shape=data.shape, order='F')
        end = min(idx+CHUNK_SIZE, data.shape[1])
        y[:,idx:end] = x[:,idx:end]

где data относится к типу np.memmap. Этот отбрасывание объекта memmap при каждом чтении предотвращает сбор массива в память и будет поддерживать очень низкое потребление памяти, если размер блока небольшой. Вероятно, это вызывает некоторые накладные расходы на ЦП, но при моей установке (MacOS) оказалось, что оно мало.

Вариант 2 - создать буфер mmap самостоятельно и дать совет по памяти

Если вы посмотрите на np.memmapИсходный код здесь, вы можете видеть, что относительно легко создать свой собственный массив с memmapped numpy. В частности, с помощью фрагмента:

mm = mmap.mmap(fid.fileno(), bytes, access=acc, offset=start)
mmap_np_array = ndarray.__new__(subtype, shape, dtype=descr, buffer=mm, offset=array_offset, order=order)

Обратите внимание на этот питон mmap экземпляр хранится как np.memmapличное _mmap атрибут.

Имея доступ к объекту mmap python и python 3.8, вы можете использовать его madviseметод, описанный здесь.

Это позволяет вам посоветовать ОС освободить память, если она доступна. Различные константы madvise описаны здесь для Linux, с некоторыми общими опциями кросса платформенных указанных.

В MADV_DONTDUMP constant выглядит многообещающе, но я не тестировал с ней потребление памяти, как для варианта 1.