Ограничить объем памяти при сохранении вывода `dask.array.map_blocks` - PullRequest
Новая
       96

Ограничить объем памяти при сохранении вывода `dask.array.map_blocks`

0 голосов
/ 29 мая 2020

Рассмотрим 2D-массив X слишком большой, чтобы поместиться в памяти - в моем случае он хранится в формате Zarr, но это не имеет значения. Я хотел бы сопоставить функцию по блокам с массивом и сохранить результат, не загружая весь массив в память. Например, 

import dask.array as da
import numpy as np

X = da.arange(10000000,
    dtype=np.int32).reshape((10,1000000)).rechunk((10,1000))

def toy_function(chunk):
    return np.mean(chunk,axis=0)

lazy_result = X.map_blocks(toy_function)

lazy_result.to_zarr("some_path")

Есть ли способ ограничить количество блоков, оцениваемых за один раз? В моем случае lazy_result[:,:1000].compute() помещается в память, но lazy_result.compute() слишком велик для памяти. Когда я пытаюсь писать в Zarr, использование памяти увеличивается, пока не достигнет максимума и не будет остановлено. Могу ли я сделать это, не прибегая к чему-то неудобному, например for i in range(1000): lazy_result[:,(i*1000):((i+1)*1000)].to_zarr('some_path'+str(i))

1 Ответ

0 голосов
/ 13 июня 2020

Я подозреваю, что ваша проблема на самом деле заключается в том, как вы создаете здесь свои исходные данные: 

X = da.arange(10000000,
    dtype=np.int32).reshape((10,1000000)).rechunk((10,1000))

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

da.ones((10, 1000000), chunks="128 MiB")
...