NumPy中广播操作的内存增长

Ces*_*sar 10 python memory numpy python-2.7

我正在使用NumPy来处理一些大型数据矩阵(大小约为50GB).我运行此代码的机器具有128GB的RAM,因此执行这种大小的简单线性操作不应该是内存方面的问题.

但是,当我在Python中计算以下代码时,我目睹了巨大的内存增长(超过100GB):

import numpy as np

# memory allocations (everything works fine)
a = np.zeros((1192953, 192, 32), dtype='f8')
b = np.zeros((1192953, 192), dtype='f8')
c = np.zeros((192, 32), dtype='f8')

a[:] = b[:, :, np.newaxis] - c[np.newaxis, :, :] # memory explodes here
Run Code Online (Sandbox Code Playgroud)

请注意,初始内存分配没有任何问题.但是,当我尝试使用广播执行减法操作时,内存增长到100GB以上.我一直认为广播会避免额外的内存分配,但现在我不确定是否总是如此.

因此,有人可以提供一些有关此内存增长发生原因的详细信息,以及如何使用更多内存有效的结构重写以下代码?

我在IPython Notebook中运行Python 2.7中的代码.

War*_*ser 6

@rth建议小批量操作是一个很好的建议.您还可以尝试使用该函数np.subtract并为其提供目标数组,以避免创建附加临时数组.我也认为你不需要索引cc[np.newaxis, :, :],因为它已经是一个3-d阵列.

而不是

a[:] = b[:, :, np.newaxis] - c[np.newaxis, :, :] # memory explodes here
Run Code Online (Sandbox Code Playgroud)

尝试

np.subtract(b[:, :, np.newaxis], c, a)
Run Code Online (Sandbox Code Playgroud)

第三个参数np.subtract是目标数组.


rth*_*rth 5

好了,您的阵列a已经占用1192953*192*32* 8 bytes/1.e9 = 58 GB了内存。

广播不会为初始数组分配额外的内存,而是

b[:, :, np.newaxis] - c[np.newaxis, :, :]
Run Code Online (Sandbox Code Playgroud)

仍保存在临时数组中。因此,在这一行中,您已分配了至少2个形状a为的数组,用于已使用的总内存>116 GB

通过一次对数组的较小子集进行操作,可以避免此问题,

CHUNK_SIZE = 100000
for idx in range(b.shape[0]/CHUNK_SIZE):
    sl = slice(idx*CHUNK_SIZE, (idx+1)*CHUNK_SIZE)
    a[sl] = b[sl, :, np.newaxis] - c[np.newaxis, :, :]
Run Code Online (Sandbox Code Playgroud)

这会稍微慢一些,但会使用更少的内存。