Numpy指定轴掩码时内存溢出问题（内存充足仍报错）

问题原因 #

当对3D数组执行arr[:, mask] = nodata时，这属于混合高级索引操作：numpy会先创建一个包含所有选中元素的临时副本，完成赋值后再将结果写回原数组。这个临时副本的大小为(第0维度长度 × mask中True的数量) × 元素字节数，即便第0维度长度为1，临时副本也会占用约7GB内存，再加上原数组、掩码以及索引计算的临时内存开销，可能触发内存溢出。

而2D数组的arr[mask] = nodata是布尔掩码的基本索引操作，numpy会直接定位原数组中的目标元素进行原地修改，无需创建大的临时副本，因此不会出现OOM错误。

解决方案 #

方案1：使用np.putmask（推荐） #

np.putmask是numpy专门用于掩码赋值的函数，支持广播机制，直接在原数组上修改元素，完全避免创建临时副本，同时兼容任意长度的第0维度：

import numpy as np

arr = np.zeros((1, 71829, 101321), dtype='uint16')
mask = np.random.randint(2, size=(71829, 101321), dtype='bool')
nodata = 0

# 原地修改，无额外内存开销
np.putmask(arr, mask, nodata)

方案2：扩展掩码维度匹配原数组 #

将掩码扩展为与原数组相同的维度，利用numpy的广播规则自动匹配第0维度，再通过布尔索引原地修改：

# 给掩码添加第0维度，shape从(H,W)变为(1,H,W)，自动广播适配原数组的(N,H,W)
mask_expanded = np.expand_dims(mask, axis=0)
arr[mask_expanded] = nodata

这两种方法的内存效率都极高，能在64GB内存环境下稳定运行，同时满足第0维度可变的需求。

内容的提问来源于stack exchange，提问作者Placibo