Numpy中3D uint8数组内存分配与释放异常问题排查

问题核心 #

在使用字典存储大量3D uint8类型Numpy数组时，首次迭代清理后内存无法释放，但2D数组（uint8/float）、3D float数组均能正常回收；第二次迭代时3D uint8数组的内存又恢复正常释放。手动触发GC、递归清理字典等操作均无效，仅通过额外分配并清理一个大float数组能临时绕过问题。

复现代码 #

import psutil
import numpy as np
import gc

# 内存检查函数
def check_mem():
    mem_dic = psutil.virtual_memory()._asdict()
    for key in mem_dic:
        mem_dic[key] = round(mem_dic[key]*1e-9)
    return mem_dic
print(f'start {check_mem()}')

toggle = False # 关闭时首次迭代3D uint8内存不释放，开启则正常

for i in range(2):
    print()
    print('iter', i)

    data = {x:np.full((10000, 10000), 1, dtype=float) for x in range(2**3)}
    print(f'make 2d float {check_mem()}')
    data.clear()
    print(f'free 2d float {check_mem()}')
    
    data = {x:np.full((3, 144, 256), 1, dtype=float) for x in range(2**12)}
    print(f'make 3d float {check_mem()}')
    data.clear()
    print(f'free 3d float {check_mem()}')
    
    data = {x:np.full((10000, 10000), 1, dtype=np.uint8) for x in range(32)}
    print(f'make 2d uint {check_mem()}')
    data.clear()
    print(f'free 2d uint {check_mem()}')
    
    data = {x:np.full((3, 144, 256), 1, dtype=np.uint8) for x in range(2**15)}
    print(f'make 3d uint {check_mem()}')
    data.clear()
    if toggle:
        data = {x:np.full((10000, 10000), 1, dtype=float) for x in range(1)}
        data.clear()
    gc.collect()
    print(f'free 3d uint {check_mem()}')

问题原因分析 #

这是Numpy内存池机制导致的行为：Numpy会为小尺寸、同类型的数组缓存内存块以提升分配效率，对于3D uint8这种固定小体积的数组，首次分配后内存池的缓存未被自动回收；而第二次迭代或分配大float数组时，会触发内存池的清理逻辑，释放缓存的内存。手动GC无法干预Numpy的内存池，因此无效。

正式解决方案 #

1. 强制清理Numpy内存池

调用Numpy内部的内存池修剪函数，直接释放缓存的内存（注意：属于内部API，未来版本可能变化，但当前可稳定解决问题）：

# 替代原toggle代码块
import numpy as np
np.core.memory._malloc_trim(0)

将原代码中toggle相关代码替换为上述内容，无需额外分配数组即可释放内存。

2. 调整数组分配策略

• 合并小尺寸数组：将多个3D uint8数组合并为一个更大的数组，减少字典中的元素数量，降低内存池缓存的规模。
• 改用np.empty初始化：np.full可能触发特定的内存分配逻辑，改用np.empty创建数组后再赋值，可避免内存池的过度缓存。

3. 升级Numpy版本

该问题可能是Numpy旧版本的已知bug，升级到最新稳定版（如1.26+），官方可能已修复内存池的回收逻辑。

4. 显式释放数组引用

确保数组的所有引用都被清除，再触发内存池清理：

# 替代原data.clear()的清理逻辑
for arr in data.values():
    del arr
data.clear()
gc.collect()
np.core.memory._malloc_trim(0)

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