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-d 数组。

所以不是

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, :, :]

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