使用numpy生成大量数据内存不足，如何边生成边写入磁盘优化？

核心问题分析 #

原代码存在两个导致内存占满的核心问题：

• 逻辑错误：遍历单样本时错误将全量x0（所有0类MNIST样本， shape为(5923, 784)）赋值给迭代变量x，每次迭代都在生成全量样本的时间步数据，额外占用数十倍内存
• 存储逻辑错误：全量1e6时间步+全量样本的数据总规模超过10PB，完全不可能在32GB内存中存放，必须采用边生成边写入磁盘的方案

优化实现（边生成边写磁盘） #

这里采用numpy.memmap实现磁盘映射数组，所有写入操作直接落地到磁盘，内存仅保留当前迭代的单步数据，同时优化冗余计算逻辑：

import numpy as np
import tensorflow as tf

def main():
    # 数据加载部分不变
    (images, labels), (_, _) = tf.keras.datasets.mnist.load_data(path="mnist.npz")
    rescaled_images = np.reshape(images, (images.shape[0], -1))
    zero_images = rescaled_images[labels == 0]
    x0 = zero_images.astype(np.float32)
    sample_num, feat_dim = x0.shape

    steps = 1000000
    mu = -2.0
    sigma = 1.5
    dt = 0.01
    # 预计算常量，避免重复计算
    mu_dt = mu * dt
    sigma_sqrt_dt = sigma * np.sqrt(dt)

    # 创建磁盘映射数组，直接对应本地npy文件，不占用内存
    # 最终数据shape为 (样本数, 时间步数+1, 特征维度)
    data_shape = (sample_num, steps + 1, feat_dim)
    fp = np.memmap('mnist_diffusion_data.npy', dtype='float32', mode='w+', shape=data_shape)

    for sample_idx in range(sample_num):
        # 写入初始时间步
        fp[sample_idx, 0] = x0[sample_idx]
        x = x0[sample_idx].copy()
        for step in range(1, steps + 1):
            # 优化冗余的ones乘法，直接用numpy广播机制
            x += mu_dt + sigma_sqrt_dt * np.random.randn(feat_dim)
            # 直接写入磁盘映射数组，不驻留内存
            fp[sample_idx, step] = x
        # 每处理完一个样本刷一次缓存到磁盘
        fp.flush()

    # 关闭映射
    del fp

if __name__ == '__main__':
    main()

额外优化建议 #

• 如果需要更高的写入效率和随机读取能力，可换成h5py操作HDF5格式，支持分块压缩，进一步减少磁盘占用
• 若不需要全量时间步数据，可根据需求间隔采样存储，大幅降低最终数据规模

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