使用multiprocessing多进程时如何安全保存TensorFlow模型

偶发模型保存后不存在、无异常抛出的问题，本质是三个多进程场景下的常见坑：

1. 多进程路径冲突：如果save_path是通过共享计数器、全局变量拼接生成，多进程并发时会出现路径重复，后写入的进程可能截断前一个进程的写入，或出现文件创建的竞态条件；如果父目录是多进程并发创建，不加exist_ok参数时会偶发目录创建失败，部分框架的save方法对父目录不存在的场景不会抛出显式异常，直接静默失败。
2. IO缓冲未刷盘：Python和操作系统默认会给文件写操作加缓冲，如果你调用save方法后主进程提前退出，子进程持有的IO缓冲会被直接丢弃，文件根本不会落到磁盘上，这个过程发生在Python解释器退出阶段，你写的try-except根本捕获不到。另外子进程的print默认也带缓冲，真的抛出异常时如果没等缓冲刷出进程就退出，控制台看不到任何报错。
3. 保存逻辑校验缺失：绝大多数框架的model.save()方法只负责发起写请求，不会等待文件完全写入磁盘就返回，没有做落盘校验的话，会出现save方法调用成功但文件实际不存在的情况。

至于跨进程返回结果报pickle错误，是因为模型对象默认绑定了大量不可序列化的运行时对象：比如线程锁、打开的文件句柄、设备上下文（CUDA流、计算图节点）、局部定义的函数等，这些对象无法被pickle序列化后跨进程传输。

• 修复偶发保存失败问题
  1. 路径生成彻底规避冲突：每个模型的save_path拼接当前进程pidos.getpid()和模型唯一标识，不要依赖跨进程共享的计数器生成文件名；保存前提前创建父目录，固定用os.makedirs(save_dir, exist_ok=True)，禁止多进程并发创建目录时不带exist_ok参数。
  2. 强制刷盘+落盘校验，替换原有保存代码为如下实现：

import os
import time

try:                       
    models[i].save(save_path)
    # 强制刷新目录项缓冲，确保文件元数据落盘
    dir_fd = os.open(os.path.abspath(os.path.dirname(save_path)), os.O_RDONLY)
    os.fsync(dir_fd)
    os.close(dir_fd)
    # 轮询确认文件真实存在且大小非0，最多等待2秒
    for _ in range(20):
        if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
            break
        time.sleep(0.1)
    else:
        raise IOError(f"Save failed: {save_path} not found on disk after write")
except Exception as E:                
    print('SAVING ERROR', E, flush=True)

3. 主进程必须等待所有子进程执行完成再退出：如果用multiprocessing.Pool，执行完任务后必须先调用pool.close()再调用pool.join()；如果是手动创建Process对象，必须对每个进程调用join()，禁止主进程提前结束运行。所有子进程中的print都加flush=True参数，避免报错信息因为缓冲被丢弃。

• 修复pickle序列化报错问题
  不要在子进程中直接返回完整模型对象，子进程完成模型训练和保存后，只返回可序列化的基础类型数据：比如模型保存的绝对路径、验证集精度、训练耗时这类字段即可。主进程拿到路径后如果需要使用模型，再在主进程中加载模型，完全绕开复杂对象跨进程序列化的问题。如果确实需要跨进程返回模型参数，只返回模型的权重字典（比如PyTorch的model.state_dict()），剥离所有绑定的运行时上下文对象即可正常序列化。

内容的提问来源于stack exchange，提问作者Роман Шарыпов