小编典典

NumPy中广播操作带来的内存增长

python

我正在使用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中运行代码。


阅读 223

收藏
2021-01-20

共1个答案

小编典典

@rth建议以较小的批次进行操作。您也可以尝试使用该函数np.subtract并将其分配给目标数组,以避免创建其他临时数组。我也认为您不需要将索引编cc[np.newaxis, :, :],因为它已经是一个3维数组。

所以代替

a[:] = b[:, :, np.newaxis] - c[np.newaxis, :, :] # memory explodes here

尝试

np.subtract(b[:, :, np.newaxis], c, a)

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

2021-01-20