Numpy memmap向量运算仍占RAM而非磁盘，出现OOM问题求助

问题描述 #

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

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显示系统内存为1.0/12.7 GB，无内存占用问题。但执行向量减法运算时，系统内存飙升至11.2/12.7 GB，最终导致内核崩溃：

result[:] = a[:] - b[:] # 仍会消耗大量内存
result.flush()

查阅np.memmap文档，明确其作用是减少内存消耗，但为何仍出现内存溢出？我猜测需要将向量减法拆分为小块处理（比如每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()

原因分析 #

a[:] - b[:]会先把整个a和b的内存映射数组强制加载到内存，执行减法后还会生成一个全新的完整临时数组。你的数组是20亿个int64元素，单个数组大小就是2e9 * 8字节 = 16GB，远超Colab的12.7GB内存上限，直接触发内存溢出。memmap本身是延迟加载，但[:]切片会一次性读入全部数据，加上运算产生的临时数组，内存直接爆掉。

正确的分块处理方案 #

你的分块思路是对的，但要注意BUFF_SIZE不能直接用字节数，得转换成元素个数——因为int64每个元素占8字节。调整后的代码如下：

import numpy as np

# 初始化memmap（如果已完成初始化可跳过此段）
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))

# 按512MB分块处理
block_size_bytes = 512 * 1024 * 1024
# 计算每个块的元素数：总字节数 / 单个元素字节数
block_size_elements = block_size_bytes // np.dtype(np.int64).itemsize

for i in range(0, result.size, block_size_elements):
    end = min(i + block_size_elements, result.size)
    # 仅加载当前块的数据到内存，运算后直接写入result
    result[i:end] = a[i:end] - b[i:end]
    result.flush()  # 可选，确保数据写入磁盘，释放临时内存

# 显式删除memmap对象，避免资源泄漏
del a, b, result

关键注意点 #

• 必须把分块的字节大小转换成元素数，否则切片范围错误，还是会加载大量数据。
• 用min()处理最后一个块，避免数组越界。
• flush()可以确保每次运算后数据写入磁盘，让numpy及时回收临时内存，不过即使不调用，numpy的内存管理也会自动处理，显式调用更稳妥。

内容的提问来源于stack exchange，提问作者Muhammad Ikhwan Perwira