Numpy memmap 在进行向量运算时仍然使用 RAM 而不是磁盘

Question

我初始化两个操作数和一个结果：

a = np.memmap('a.mem', mode='w+', dtype=np.int64, shape=(2*1024*1024*1024))
b = np.memmap('b.mem', mode='w+', dtype=np.int64, shape=(2*1024*1024*1024))
result = np.memmap('result.mem', mode='w+', dtype=np.int64, shape=(2*1024*1024*1024))

在这样的空闲状态下，Google Colab 报告的系统 RAM 仍然

1.0/12.7 GB

，这很好，还没有 RAM 活动。但是，进行矢量减法等矢量操作时，报告的系统内存增加到几乎最大峰值，即

11.2/12.7 GB

，最终运行时内核崩溃：

result[:] = a[:] - b[:] # This still consume memory
result.flush()

我已经多次阅读

np.memmap

文档，据说

memmap

的目的应该是减少内存消耗，但为什么我仍然遇到

Out Of Memory

错误？

我怀疑，向量减法必须缓冲成小块，例如每个

512MB

缓冲存储器。但我不知道语法是什么： 也许我的意思是这样的：

BUFF_SIZE = 512 * 1024 * 1024
for i in range(0, result.size, BUFF_SIZE):
  result[i:i+BUFF_SIZE] = a[i:i+BUFF_SIZE] - b[i:i+BUFF_SIZE]
  result.flush()

Answer 1

result[:] = a[:] - b[:]

并不意味着“将减法结果写入

result

”。它的意思是“将减法结果写入一个新数组，然后将该数组的内容复制到

result

”。您正在尝试在中间创建一个 16 GiB 临时阵列。

要将输出直接写入

result

，可以使用

out

ufunc的

numpy.subtract

参数：

numpy.subtract(a, b, out=c)