如何分批保存压缩Numpy数组且无需同时加载全部至内存？

我们可以通过逐次写入的方式分批保存多个Numpy数组，不用一次性把所有数组都加载到内存里，示例代码如下：

with open('test.npy', 'wb') as f:
    A = compute_my_np_array(1)
    np.save(f, A)
    # 可以手动执行 del A 释放内存，也能靠后续赋值自动覆盖释放

    A = compute_my_np_array(2)
    np.save(f, A)

但这种方式存出来的数组是未压缩的，而官方的np.savez_compressed必须把所有数组都加载到内存才能完成压缩保存：

A = compute_my_np_array(1)
B = compute_my_np_array(2)
np.savez_compressed('test.npz', A=A, B=B)

核心问题：如何分批保存压缩后的Numpy数组，且无需同时加载全部数组到内存？

解决方案：用zipfile手动分批压缩保存 #

.npz文件本质就是个ZIP压缩包，里面装着多个.npy格式的文件。所以我们可以直接用Python标准库zipfile来逐个生成压缩后的数组文件，实现分批处理：

import numpy as np
import zipfile

def compute_my_np_array(idx):
    # 这里是示例数组生成逻辑，替换成你自己的计算代码就行
    return np.random.rand(1000, 1000)

with zipfile.ZipFile('test_compressed.npz', 'w', zipfile.ZIP_DEFLATED) as zf:
    # 处理第一个数组
    arr = compute_my_np_array(1)
    # 把数组写入ZIP包里的对应文件
    with zf.open('arr_1.npy', 'w') as f:
        np.save(f, arr)
    # 内存吃紧的话可以手动释放内存
    del arr

    # 处理第二个数组
    arr = compute_my_np_array(2)
    with zf.open('arr_2.npy', 'w') as f:
        np.save(f, arr)
    del arr

    # 还可以继续添加更多数组...

加载分批保存的压缩数组 #

加载的时候也能按需读取，不用一次性把所有数组都加载进来：

with zipfile.ZipFile('test_compressed.npz', 'r') as zf:
    # 读取第一个数组
    with zf.open('arr_1.npy', 'r') as f:
        arr1 = np.load(f)
    print(arr1.shape)

    # 读取第二个数组
    with zf.open('arr_2.npy', 'r') as f:
        arr2 = np.load(f)
    print(arr2.shape)

补充说明 #

• 这种方式完全不需要同时加载所有数组到内存，处理完一个就能释放内存再处理下一个
• 使用zipfile.ZIP_DEFLATED启用压缩，和np.savez_compressed的压缩效果一致
• 可以自定义ZIP包内的文件名（比如arr_1.npy、arr_2.npy），方便后续识别读取

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