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

Question

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

Answer 1

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

所以代替

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

尝试

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

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