如何同时训练多个独立的TensorFlow神经网络？

现有方案的潜在问题 #

你当前用multiprocessing的思路方向没问题，但代码存在几个关键问题，可能导致训练异常或资源浪费：

1. 模型与资源的跨进程冲突：TensorFlow的模型、会话资源是绑定进程的，主进程创建的模型直接在子进程中使用，会引发序列化（Windows下pickle报错）或资源竞争问题（比如多个进程抢占同一GPU内存，导致OOM）。
2. 资源分配失控：如果同时开启大量进程，会导致CPU上下文切换开销剧增，或GPU内存被占满，反而降低训练效率。
3. 扩展性差：手动为每个模型写单独的训练函数，当模型数量增多时（比如几十上百个），代码会变得极度冗余。

优化后的实现方法 #

方法1：改进多进程写法——每个进程独立创建模型与加载数据 #

核心思路是让每个子进程完全独立，在进程内部创建模型、加载数据集，避免跨进程的资源共享问题。同时用通用训练函数替代单个模型的训练函数，提升扩展性。

示例代码：

import multiprocessing
import tensorflow as tf

def train_single_model(layers, dataset_path, epochs, batch_size, compile_params):
    # 子进程内独立创建模型
    model = tf.keras.Sequential(layers)
    model.compile(**compile_params)
    # 独立加载数据集（建议提前用tf.data保存，避免重复预处理）
    dataset = tf.data.experimental.load(dataset_path)
    # 执行训练
    model.fit(dataset, epochs=epochs, batch_size=batch_size)
    # 保存训练好的模型
    model.save(f"./trained_model_{compile_params['optimizer']}_test.h5")

if __name__ == '__main__':
    epochs = 100
    batch_size = 32

    # 批量定义模型配置（可从文件读取，支持扩展到大量模型）
    model_tasks = [
        {
            "layers": [tf.keras.layers.Dense(64, activation='relu'), tf.keras.layers.Dense(1)],
            "dataset_path": "./dataset_A",
            "compile_params": {"optimizer": "adam", "loss": "mean_squared_error", "metrics": ["accuracy"]}
        },
        {
            "layers": [tf.keras.layers.Dense(128, activation='relu'), tf.keras.layers.Dense(10, activation='softmax')],
            "dataset_path": "./dataset_B",
            "compile_params": {"optimizer": "sgd", "loss": "categorical_crossentropy", "metrics": ["accuracy"]}
        }
    ]

    # 控制并发进程数（建议设为CPU核心数或GPU数量，避免资源过载）
    max_workers = multiprocessing.cpu_count() // 2
    with multiprocessing.Pool(max_workers) as pool:
        # 批量提交训练任务
        for task in model_tasks:
            pool.apply_async(
                train_single_model,
                args=(task["layers"], task["dataset_path"], epochs, batch_size, task["compile_params"])
            )
        pool.close()
        pool.join()

方法2：用Ray框架管理大规模训练任务 #

如果需要训练的模型数量非常多（上百个），或者需要跨机器扩展，multiprocessing的资源调度能力不足，推荐用Ray——专为机器学习设计的分布式框架，能自动管理CPU/GPU资源，支持任务队列和弹性扩展。

示例代码（需先安装pip install ray）：

import ray
import tensorflow as tf

# 初始化Ray
ray.init()

# 定义远程训练函数，指定资源占用（比如每个任务用1个CPU，或0.5个GPU）
@ray.remote(num_cpus=1)
def train_single_model(layers, dataset_path, epochs, batch_size, compile_params):
    model = tf.keras.Sequential(layers)
    model.compile(**compile_params)
    dataset = tf.data.experimental.load(dataset_path)
    model.fit(dataset, epochs=epochs, batch_size=batch_size)
    model.save(f"./trained_model_{compile_params['optimizer']}_ray.h5")
    return model.evaluate(dataset)

if __name__ == '__main__':
    epochs = 100
    batch_size = 32
    model_tasks = [
        # 和方法1相同的任务配置
    ]

    # 批量提交任务，返回异步结果对象
    futures = [
        train_single_model.remote(
            task["layers"], task["dataset_path"], epochs, batch_size, task["compile_params"]
        ) for task in model_tasks
    ]

    # 等待所有任务完成并获取结果
    results = ray.get(futures)
    print("所有模型训练结果：", results)
    ray.shutdown()

方法3：GPU场景下的资源隔离 #

如果用GPU训练，需要为每个进程分配独立的GPU资源，避免内存冲突。可以在训练函数开头添加GPU指定逻辑：

def train_single_model(...):
    # 获取可用GPU列表，按进程ID分配不同GPU
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        # 假设进程ID从0开始，循环分配GPU
        process_id = multiprocessing.current_process().pid % len(gpus)
        tf.config.set_visible_devices(gpus[process_id], 'GPU')
        tf.config.experimental.set_memory_growth(gpus[process_id], True)  # 按需分配内存

关键注意事项 #

• 数据集预处理：建议提前将每个数据集用tf.data.experimental.save保存，避免子进程重复执行预处理逻辑，节省时间。
• 并发数控制：不要开启超过硬件承载能力的进程数——CPU场景下建议设为核心数的1-2倍，GPU场景下最多等于GPU数量。
• 模型保存：每个模型训练完成后要单独保存，避免覆盖或丢失结果。

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