我正在使用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
请注意,初始内存分配已完成,没有任何问题。但是,当我尝试通过广播执行减法运算时,内存会增加到100GB以上。我一直认为广播会避免进行额外的内存分配,但是现在我不确定是否总是这样。
这样,有人可以提供一些有关为何这种内存增长的原因的详细信息,以及如何使用内存效率更高的结构重写以下代码吗?
我正在IPython Notebook中的Python 2.7中运行代码。
@rth建议以较小的批次进行操作。您也可以尝试使用该函数np.subtract并将其分配给目标数组,以避免创建其他临时数组。我也认为您不需要将索引编c为c[np.newaxis, :, :],因为它已经是一个3维数组。
np.subtract
c
c[np.newaxis, :, :]
所以代替
a[:] = b[:, :, np.newaxis] - c[np.newaxis, :, :] # memory explodes here
尝试
np.subtract(b[:, :, np.newaxis], c, a)
的第三个参数np.subtract是目标数组。