Python模拟任务中大规模数据的高效累积方案咨询

核心需求梳理 #

• 生成数据总量约10GB，无法全量缓存到内存
• 当前用CPU多进程（12核）运行，单轮耗时约1小时，需避免写入操作大幅增加耗时
• 后续需支持GPU计算、更高配置服务器迁移

推荐方案 #

1. 用PyTorch原生二进制格式（.pt/.pth）分批次写入 #

这是最贴合当前PyTorch技术栈的低损耗方案：

• 操作步骤：
  1. 每个进程的simulate函数生成x、y后，直接调用torch.save({"x": x, "y": y}, f"proc_{os.getpid()}_batch_{batch_idx}.pt")，按进程+批次序号命名独立二进制文件
  2. 后续读取时直接用torch.load加载，可直接转换为GPU张量
• 优势：
  • 二进制序列化/反序列化速度远快于文本格式，几乎不增加额外耗时
  • 完美保留张量的 dtype、形状信息，无需手动解析转换
  • 后续迁移GPU时，加载后直接调用.to("cuda")即可无缝衔接
• 注意事项：
  • 用进程ID+批次号命名，避免多进程写入时的文件冲突
  • 可按子目录分组存放批次文件，避免单目录文件数量过多

2. 使用HDF5格式（h5py库）并行写入 #

适合需要结构化存储、后续需做数据分析的场景：

• 操作步骤：
  1. 初始化HDF5文件，创建可扩展的数据集：

     import h5py
     import numpy as np
     
     with h5py.File("sim_data.h5", "w", libver="latest") as h5f:
         h5f.create_dataset("x", shape=(0, LENGTH, WIDTH), maxshape=(None, LENGTH, WIDTH), dtype=np.int6)
         h5f.create_dataset("y", shape=(0,7), maxshape=(None,7), dtype=np.int16)
     

  2. 每个进程生成批次数据后，开启SWMR（单写多读）模式安全写入：

     with h5py.File("sim_data.h5", "a", swmr=True) as h5f:
         current_len = h5f["x"].shape[0]
         # 扩展数据集并写入x
         h5f["x"].resize(current_len + batch_size, axis=0)
         h5f["x"][current_len:] = x.numpy()
         # 扩展数据集并写入y
         h5f["y"].resize(current_len + batch_size, axis=0)
         h5f["y"][current_len:] = y.numpy()
     

• 优势：
  • 支持随机读写和数据分片，后续可按需加载部分数据
  • 磁盘占用与二进制格式接近，性能损耗小
  • 兼容Python数据分析生态（Pandas、NumPy），也可被PyTorch直接读取
• 注意事项：
  • 必须开启swmr=True确保多进程写入的线程安全
  • 对齐当前batch_size=50的批次大小写入，减少IO频次

3. 多进程异步写入队列（解耦计算与IO） #

若担心写入操作阻塞计算进程，可通过队列实现并行：

• 操作步骤：
  1. 主进程创建multiprocessing.Queue（设置合理maxsize，比如10），用于接收计算好的批次数据
  2. 启动1-2个独立IO进程，负责从队列取数据并写入文件（采用上述.pt或HDF5方式）
  3. 计算进程仅负责生成x、y，将数据丢入队列后立刻返回继续计算
• 优势：
  • 计算与IO操作并行，避免写入拖慢计算进程的CPU利用率
  • 适配IO性能较弱的机器，最大化利用计算资源
• 注意事项：
  • 队列maxsize不宜过大，避免内存堆积过多未写入数据
  • IO进程需做好异常捕获，防止队列数据丢失

避坑提醒 #

• 禁止使用CSV/TXT存储：文本格式序列化速度慢，磁盘占用会膨胀3倍以上，后续解析还需额外做类型转换，完全不适合大规模张量数据
• 不要用pickle序列化大列表：会触发内存溢出问题，且序列化速度远不如PyTorch原生格式
• 后续迁移GPU时优先选择.pt格式，加载后可直接转GPU张量，无需中间转换步骤

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