Как обработать большой тензор numpy?

Question

Имеется тензор, сохранённый в файле .npy, который читаю так:

x = np.load('x.npy')

Задача состоит в том, чтобы выполнить с ним математические преобразования и сохранить в другой файл.

x = process(x)
np.save('x_processed.npy', x)

Исходный тензор помещается в оперативную память, возвращаемое функцией process значение - нет.

Можно преобразовывать данные по частям: применять функцию process к срезу тензора x, но формат .npy не поддерживает добавление (append).

Я попробовал применить numpy.memmap таким образом:

x = np.load('x.npy', mmap_mode='r+')

Ожидаемое поведение - тензор не будет загружен в оперативную память, вместо этого будет создано отображение файла в оперативную память. Если в дальнейшем редактировать данные в тензоре x, то они будут редактироваться именно в файле.

Но функция process построена именно так, что она возвращает вновь созданный тензор, поэтому он создаётся в оперативной памяти и полностью её съедает.

Всё-таки есть какой-то способ обрабатывать исходный тензор по частям (срезам)?

Пусть имеются два файла, полученные путём обработки среза функцией process, которые затем загружаются посредством mmap:

x1 = np.load('x1.npy', mmap_mode='r')
x2 = np.load('x2.npy', mmap_mode='r')

И файл, куда надо записать результат:

x_targ = np.memmap(filename, dtype='float32', mode='w+', shape=full_tensor_shape)

тогда можно ли делать так:

x_targ[0:x1.shape[0], :, :, :] = x1
x_targ[0:x2.shape[0], :, :, :] = x2

Или как их объединить в один файл?

Answer 1

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

x = np.load('x.npy', mmap_mode='r')
x_out_shape = (x.shape[0], A, B, C)
x_out = np.memmap('x_out.npy', dtype='uint8', mode='w+', shape=x_out_shape)
particle = 1000

filePos = 0
while True:
    # Если filePos + particle больше размера тензора, 
    # то автоматически подставится верхняя граница
    x_out[filePos:filePos + particle, :, :, :] = process(x[filePos:filePos + particle, :, :, :])
    filePos += particle
    if filePos >= x_out_shape[0]:
        break
x_out.flush()

БЛОГ НА HUSL

Как обработать большой тензор numpy?

Ответы (1 шт):