在Keras（TensorFlow后端）CPU上并行训练多架构小型神经网络的方法

嘿，这个场景我太熟悉了——当你有一堆小模型要训，默认TF的多线程反而让核心利用率上不去，还浪费资源。下面给你几个经过实践验证的方案，完美适配你的需求：

方法1：多进程绑定单个CPU核心（最推荐） #

因为TensorFlow在单进程内会自动抢占多个核心，所以给每个模型分配独立进程并绑定到单个CPU核心是最直接的思路，能彻底隔离每个模型的计算资源，避免互相干扰。

实现步骤： #

1. 用Python的multiprocessing创建进程池，每个进程负责一个模型的训练。
2. 在每个进程内部，设置CPU亲和性，把进程绑定到指定核心。
3. 每个进程单独导入TensorFlow/Keras，避免进程间的TF状态冲突。

代码示例： #

import multiprocessing
import os
import psutil  # 跨平台CPU亲和性设置，需要先安装：pip install psutil

def train_model_on_core(model_config, core_id):
    # 每个进程单独导入TF/Keras，防止全局状态冲突
    import tensorflow as tf
    from tensorflow import keras

    # 绑定当前进程到指定核心
    p = psutil.Process(os.getpid())
    p.cpu_affinity([core_id])  # 传入单个核心的ID，比如0、1、2...

    # 禁用TF的自动多线程，强制单核心运行
    tf.config.threading.set_intra_op_parallelism_threads(1)
    tf.config.threading.set_inter_op_parallelism_threads(1)

    # 根据model_config构建异构模型（这里示例用简单模型，实际替换成你的模型构建逻辑）
    model = keras.Sequential(model_config["layers"])
    model.compile(optimizer=model_config["optimizer"], loss=model_config["loss"])

    # 训练模型（替换成你的数据和训练参数）
    model.fit(model_config["x_train"], model_config["y_train"], epochs=10, batch_size=32)

if __name__ == "__main__":
    # 假设你有10个不同架构的模型配置
    model_configs = [
        {"layers": [keras.layers.Dense(32, activation='relu', input_shape=(10,)), keras.layers.Dense(1)],
         "optimizer": "adam", "loss": "mse", "x_train": ..., "y_train": ...},
        # 其他模型配置...
    ]

    # 获取CPU核心数，假设你要训的模型数不超过核心数
    num_cores = multiprocessing.cpu_count()
    processes = []

    for i, config in enumerate(model_configs):
        core_id = i % num_cores  # 循环分配核心
        p = multiprocessing.Process(target=train_model_on_core, args=(config, core_id))
        processes.append(p)
        p.start()

    # 等待所有进程完成
    for p in processes:
        p.join()

优点： #

• 完全隔离每个模型的计算资源，核心利用率能拉满
• 天然支持异构模型，每个进程的模型架构可以完全不同
• 跨平台兼容（Windows/Linux/macOS）

方法2：用TF的OneDeviceStrategy指定核心 #

如果你更倾向于用TensorFlow原生的分布式策略，可以给每个模型分配一个单独的CPU设备，强制模型在指定核心上运行。

代码示例： #

import tensorflow as tf
from tensorflow import keras

def train_model_on_device(model_config, device):
    # 创建单设备策略
    strategy = tf.distribute.OneDeviceStrategy(device=device)

    with strategy.scope():
        # 在策略作用域内构建和编译模型
        model = keras.Sequential(model_config["layers"])
        model.compile(optimizer=model_config["optimizer"], loss=model_config["loss"])

    # 训练模型
    model.fit(model_config["x_train"], model_config["y_train"], epochs=10, batch_size=32)

# 假设你有4个核心，分配给4个模型
devices = ["/CPU:0", "/CPU:1", "/CPU:2", "/CPU:3"]
model_configs = [...]  # 你的异构模型配置列表

# 用多线程或多进程启动训练（注意：单线程下多个策略会串行，所以还是要配合多进程）
import threading
threads = []
for config, device in zip(model_configs, devices):
    t = threading.Thread(target=train_model_on_device, args=(config, device))
    threads.append(t)
    t.start()

for t in threads:
    t.join()

注意点： #

• 单独用多线程的话，TF的全局线程池可能还是会互相抢占资源，所以最好还是配合方法1的多进程+核心绑定
• 设备名的格式是"/CPU:N"，N从0开始对应不同的核心

方法3：全局限制TF线程数后并行训练 #

如果你不想用多进程，可以先全局设置TF只使用1个线程，然后用多线程来跑多个模型。这种方式更轻量，但隔离性稍差。

代码示例： #

import tensorflow as tf
from tensorflow import keras
import threading

# 全局设置TF单线程运行
tf.config.threading.set_intra_op_parallelism_threads(1)
tf.config.threading.set_inter_op_parallelism_threads(1)

def train_model(model_config):
    model = keras.Sequential(model_config["layers"])
    model.compile(optimizer=model_config["optimizer"], loss=model_config["loss"])
    model.fit(model_config["x_train"], model_config["y_train"], epochs=10, batch_size=32)

model_configs = [...]  # 你的异构模型配置列表

threads = []
for config in model_configs:
    t = threading.Thread(target=train_model, args=(config,))
    threads.append(t)
    t.start()

for t in threads:
    t.join()

缺点： #

• 线程间共享TF的全局状态，可能会有隐蔽的资源冲突
• 核心绑定需要依赖操作系统的调度，不如多进程精准

关键注意事项 #

• 进程隔离优先：TensorFlow的全局状态（比如默认图、线程池）在多进程下是完全隔离的，所以尽量用multiprocessing而不是threading，避免奇怪的bug
• 内存监控：每个进程都会加载TensorFlow和模型，所以如果模型数量很多，要注意内存是否足够（小模型的话问题不大）
• 核心数匹配：尽量让模型数量和CPU核心数一致，避免核心过载或闲置

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