如何高效求和大型(100,1024,1024) .npz矩阵？求更优存储格式

一、解决内存不足的求和方案 #

你的原代码会一次性把所有矩阵加载到内存再求和，每个(100,1024,1024)的float32矩阵占约400MB，10个就占4GB，float64直接翻倍，很容易触发内存不足。改成逐文件加载累加，只保留总和矩阵在内存里，就能大幅降低内存占用：

import numpy as np

# 初始化总和矩阵：用第一个文件的shape和数据类型创建全零矩阵
first_mat = np.load(files_list[0])
summed_matrix = np.zeros_like(first_mat)
del first_mat  # 立即释放临时内存

# 循环处理每个文件，累加后释放单个矩阵内存
for file_path in files_list:
    current_mat = np.load(file_path)
    summed_matrix += current_mat
    del current_mat  # 让GC及时回收单个矩阵的内存

如果你的.npz文件里包含多个数组（不是单矩阵），记得指定键名，比如np.load(file_path)['your_array_key']。

二、更优的存储格式建议 #

针对(100,1024,1024)这类大型多维数组，以下几种格式比.npz更适合：

1. HDF5（h5py库） #

HDF5支持分块存储和按需读取，不需要把整个数组加载到内存就能操作。你可以把多个矩阵存在同一个HDF5文件的不同数据集里，求和时逐数据集读取累加，甚至直接在磁盘上做部分运算：

import h5py

# 写入示例：把多个矩阵存到HDF5
with h5py.File('matrices.h5', 'w') as f:
    for i, file_path in enumerate(files_list):
        mat = np.load(file_path)
        f.create_dataset(f'mat_{i}', data=mat, compression='gzip')  # 启用压缩

# 求和示例：逐数据集读取累加
summed_matrix = None
with h5py.File('matrices.h5', 'r') as f:
    for ds_name in f.keys():
        ds = f[ds_name]
        if summed_matrix is None:
            summed_matrix = np.zeros_like(ds)
        summed_matrix += ds[:]  # 按需读取整个数据集（也可以分块读）

2. Zarr #

Zarr和HDF5功能类似，但API更简洁，基于ZIP压缩，还支持云存储，适合并行处理。它的内存友好性极佳，支持分块读写：

import zarr

# 写入示例
store = zarr.DirectoryStore('matrices.zarr')
root = zarr.group(store=store)
for i, file_path in enumerate(files_list):
    mat = np.load(file_path)
    root.create_dataset(f'mat_{i}', data=mat, compression='zstd')

# 求和示例
summed_matrix = None
root = zarr.open('matrices.zarr', mode='r')
for ds_name in root:
    ds = root[ds_name]
    if summed_matrix is None:
        summed_matrix = np.zeros_like(ds)
    summed_matrix += ds[:]

3. 优化现有.npz存储 #

如果不想换格式，也可以做两点优化：

• 用np.savez_compressed替代普通np.savez，它会用更高压缩比存储，减少磁盘占用和加载时的内存开销。
• 降低数据类型：如果精度允许，把float64转成float32，int64转成int32，能直接减少一半的内存和存储占用，比如mat = mat.astype(np.float32)。

4. 内存映射文件（numpy.memmap） #

把数组存储为二进制文件，通过内存映射的方式访问，不需要一次性加载整个数组到内存。适合超大型数组的累加操作：

# 创建内存映射的总和矩阵
summed_mmap = np.memmap('summed_matrix.dat', dtype=np.float32, mode='w+', shape=(100,1024,1024))

# 逐文件累加
for file_path in files_list:
    mat = np.load(file_path).astype(np.float32)
    summed_mmap += mat
    del mat

# 写入磁盘并关闭
summed_mmap.flush()
del summed_mmap

# 后续读取
summed_matrix = np.memmap('summed_matrix.dat', dtype=np.float32, mode='r', shape=(100,1024,1024))

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