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, :, :]，因为它已经是一个三维数组。

所以而不是

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

试

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

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

Answer 2

好吧，你的数组a已经占用了1192953*192*32* 8 bytes/1.e9 = 58 GB个内存。

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

的结果

b[:, :, np.newaxis] - c[np.newaxis, :, :]

仍保存在临时数组中。因此，在此行中，您已分配至少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, :, :]

这将稍微慢一些，但使用的内存要少得多。

NumPy中广播操作的内存增长

2 个答案: