TensorFlow Keras模型无法在多进程中执行预测的问题排查

我来帮你拆解这个问题——TensorFlow和Python多进程的交互确实有不少坑，你遇到的卡住问题主要是两个核心原因导致的：

问题成因 #

1. 进程继承的TensorFlow状态混乱：Python的multiprocessing.Process默认在Unix系统下用fork方式创建子进程，这种方式会直接复制父进程的内存空间，包括TensorFlow已经初始化的全局上下文、计算图或者资源锁。子进程继承这些状态后，会和父进程或者其他子进程争夺共享资源，最终导致死锁，进程永远卡在那里。
2. 子进程重复初始化模型的资源冲突：你的代码里每个子进程都重新创建并初始化Sequential模型，TensorFlow初始化模型时会分配底层的CPU/GPU资源，多进程同时做这件事会触发资源竞争，同样会导致进程挂起。

你之前尝试的tf.device('/cpu:0')没用，是因为问题根本不在设备选择上，而是进程间的状态和资源冲突。

可行的解决方案 #

下面给你几种不同场景下的解决办法，你可以根据自己的需求选择：

方案一：在子进程内重置TensorFlow状态 #

在每个子进程的预测函数开头，先清理TensorFlow的全局状态，确保子进程从头开始初始化，避免继承父进程的混乱状态：

import multiprocessing as mp
import tensorflow as tf
import numpy as np

def predict(data):
    # 清理之前的TensorFlow上下文，重置状态
    tf.keras.backend.clear_session()
    # 可选：设置随机种子保证结果可复现
    tf.random.set_seed(42)
    a = tf.keras.Sequential([tf.keras.layers.Dense(4, input_shape=(16,))])
    result = a.predict(data)
    # 用完再清理一次，释放资源
    tf.keras.backend.clear_session()
    return result

fake_data = np.zeros((100, 16))

# 单进程测试正常
for i in range(4):
    print(predict(fake_data).shape)

# 多进程测试
processes = []
for i in range(4):
    p = mp.Process(target=predict, args=(fake_data,))
    p.start()
    processes.append(p)
for p in processes:
    p.join()

方案二：改用spawn方式创建进程 #

spawn是另一种进程创建方式，子进程会从头启动Python解释器，完全不继承父进程的任何状态（包括TensorFlow的上下文），从根源上避免状态冲突：

import multiprocessing as mp
import tensorflow as tf
import numpy as np

def predict(data):
    a = tf.keras.Sequential([tf.keras.layers.Dense(4, input_shape=(16,))])
    return a.predict(data)

if __name__ == "__main__":
    fake_data = np.zeros((100, 16))
    
    # 单进程测试
    for i in range(4):
        print(predict(fake_data).shape)
    
    # 显式使用spawn上下文创建进程
    ctx = mp.get_context('spawn')
    processes = []
    for i in range(4):
        p = ctx.Process(target=predict, args=(fake_data,))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

注意：Windows系统默认用spawn，但Unix/Linux/macOS需要手动指定。这种方式启动进程会慢一点，因为要重新加载所有模块，但能彻底解决TensorFlow的多进程状态问题。

方案三：用进程池共享模型（更高效） #

如果你的模型是固定的，不要在每个子进程里重复创建模型——这太浪费资源了。可以用multiprocessing.Pool，在每个子进程启动时只初始化一次模型，之后重复使用：

import multiprocessing as mp
import tensorflow as tf
import numpy as np

# 全局变量存储子进程内的模型
global_model = None

def init_worker():
    # 每个子进程启动时初始化一次模型
    global global_model
    tf.keras.backend.clear_session()
    global_model = tf.keras.Sequential([tf.keras.layers.Dense(4, input_shape=(16,))])

def predict(data):
    # 直接用已经初始化好的模型做预测
    return global_model.predict(data)

if __name__ == "__main__":
    fake_data = np.zeros((100, 16))
    
    # 单进程测试
    init_worker()
    for i in range(4):
        print(predict(fake_data).shape)
    
    # 使用进程池，每个子进程只初始化一次模型
    with mp.Pool(processes=4, initializer=init_worker) as pool:
        results = pool.map(predict, [fake_data]*4)
        for res in results:
            print(res.shape)

这种方式效率最高，因为模型只初始化4次（对应4个进程），而不是每次预测都重新创建，适合需要多次预测的场景。

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