如何快速读写3D Numpy数组？循环生成(W,N)数组的高效存储方案

咱们先理清核心问题：你要处理的规模是M≈1e5-1e6个(W=500,N=100)的数组，先算笔内存账——每个数组是500*100=5e4个元素，如果用float64类型，每个数组占400KB，1e6个就是400GB，这远远超过普通机器的内存容量，所以你之前尝试的“一次性预分配大numpy数组”根本不可能在内存里放下；而np.append/vstack慢到离谱的原因是：这些操作每次都会创建新数组并复制原有数据，时间复杂度是O(M²)，对于1e6的规模完全不可行。

下面分两种场景给你最优方案：

一、内存能装下的小M场景（比如M≤1e5，内存≥40GB） #

如果你的M比较小，内存足够容纳整个数据集，预分配numpy数组直接赋值是最快的方式——一次性申请好内存块，循环里直接写入，没有任何复制开销：

import numpy as np

M = 100000
W = 500
N = 100

# 预分配数组，指定dtype避免自动转换开销
result = np.empty((M, W, N), dtype=np.float64)

for i in range(M):
    # 生成你的(W,N)数组
    current_arr = generate_wn_array()
    # 直接赋值，O(1)内存写入操作
    result[i] = current_arr

这种方式的速度接近内存写入的极限，是纯内存操作里最快的。

二、内存装不下的大M场景（M≥1e5） #

这时候必须用“磁盘+内存映射”的方案，把数据存在磁盘上，像访问内存数组一样操作，只加载需要的部分到内存。推荐以下几种方案：

1. 优先用numpy原生的memmap（最简单，无额外依赖） #

numpy的内存映射文件直接把数组存在磁盘，接口和普通numpy数组完全一致，顺序写入速度极快，随机访问也很方便：

import numpy as np

M = 10**6
W = 500
N = 100
dtype = np.float64

# 创建内存映射文件，mode='w+'表示可读写
mmap_arr = np.memmap('large_data.npy', dtype=dtype, mode='w+', shape=(M, W, N))

for i in range(M):
    current_arr = generate_wn_array()
    # 直接写入磁盘对应位置，无内存复制
    mmap_arr[i] = current_arr

# 写完后删除引用，确保数据写入磁盘
del mmap_arr

# 后续读取时
mmap_arr = np.memmap('large_data.npy', dtype=dtype, mode='r', shape=(M, W, N))
# 访问第k个(W,N)数组，只会加载对应块到内存
print(mmap_arr[k])

2. 需要并行生成/复杂计算？用Dask数组 #

如果你的数组生成过程可以并行，或者后续要做数据分析、统计等操作，Dask是绝佳选择——它自动把大数组分成小块，支持并行计算，内存占用可控：

import dask.array as da
from dask.diagnostics import ProgressBar

M = 10**6
W = 500
N = 100
chunk_size = 1000  # 每个块1000个(W,N)数组，根据你的内存调整

# 定义生成单块数组的函数
def generate_chunk(_):
    chunk = np.empty((chunk_size, W, N), dtype=np.float64)
    for i in range(chunk_size):
        chunk[i] = generate_wn_array()
    return chunk

# 创建Dask数组，自动分块
dask_arr = da.map_blocks(
    generate_chunk,
    chunks=(chunk_size, W, N),
    dtype=np.float64,
    shape=(M, W, N)
)

# 写入磁盘（用zarr格式，支持分块存储和压缩）
with ProgressBar():
    dask_arr.to_zarr('large_data.zarr', overwrite=True)

# 读取并访问
dask_arr = da.from_zarr('large_data.zarr')
# 计算第k个元素，Dask自动加载对应块
print(dask_arr[k].compute())

3. 长期存储/多语言访问？用h5py+HDF5 #

HDF5是工业级的大规模数据存储格式，支持分块、压缩、随机访问，而且可以用Python、C++、MATLAB等多种语言读取，适合长期归档：

import h5py

M = 10**6
W = 500
N = 100

# 创建HDF5文件，指定分块大小
with h5py.File('large_data.h5', 'w') as f:
    dset = f.create_dataset(
        'data',
        shape=(M, W, N),
        dtype='f8',
        chunks=(1000, W, N)  # 分块大小和内存匹配
    )
    for i in range(M):
        dset[i] = generate_wn_array()

# 读取时
with h5py.File('large_data.h5', 'r') as f:
    dset = f['data']
    # 直接访问第k个(W,N)数组
    print(dset[k])

总结 #

• 小M内存足够：预分配numpy数组直接赋值，速度最快。
• 大M内存不足：优先用numpy.memmap（简单无依赖）；需要并行/复杂计算用Dask；长期存储用h5py+HDF5。
• 完全不需要用np.append这类操作，它们的时间复杂度对于大规模数据来说是灾难级的。

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