Python中TensorFlow LSTM模型多进程预测报错挂起问题排查

问题根源 #

TensorFlow的运行时状态和Python多进程的fork机制冲突，父进程的TF上下文会被子进程继承，导致模型加载、指标解析出现异常；另外模型加载时的指标序列化在跨进程场景下容易出问题，哪怕你没自定义指标也会踩坑。

解决办法 #

1. 强制用spawn模式启动多进程 #

Unix系统默认的fork模式会继承父进程的TF状态，改用spawn模式能让每个子进程完全独立初始化TF环境，避免状态混乱。

import multiprocessing as mp

if __name__ == "__main__":
    # 强制设置启动方式为spawn
    mp.set_start_method('spawn', force=True)
    # 初始化进程池
    pool = mp.Pool(processes=4)
    # 批量执行预测
    results = pool.map(your_predict_function, your_data_list)
    pool.close()
    pool.join()

2. 子进程内单独初始化TF并加载模型 #

别在父进程里提前加载模型，每个子进程自己导入TF、加载模型，确保环境完全独立。加载时可以显式指定内置指标，避免解析错误。

def predict_worker(data):
    # 子进程内单独导入TensorFlow和模型加载工具
    import tensorflow as tf
    from tensorflow.keras.models import load_model

    # 显式指定内置指标，避免跨进程解析失败
    model = load_model('your_lstm_model.h5', custom_objects={'accuracy': tf.keras.metrics.CategoricalAccuracy})
    # 执行预测
    result = model.predict(data)
    return result

3. 父进程别提前碰TF相关内容 #

所有TF的初始化、模型加载代码都要放在if __name__ == "__main__"里面或者子进程函数里，防止子进程继承无效的TF上下文。

4. 换用“保存权重+重建模型”的方式 #

如果还是有指标问题，可以只保存模型权重，子进程里重新构建模型结构再加载权重，绕开指标序列化的坑。

# 父进程中保存权重（训练完之后）
model.save_weights('lstm_weights.h5')

# 子进程中重建模型并加载权重
def predict_worker(data):
    import tensorflow as tf
    # 完全复刻训练时的模型结构
    def build_lstm_model():
        model = tf.keras.Sequential([
            tf.keras.layers.LSTM(64, input_shape=(你的时间步长, 特征数)),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        return model
    
    model = build_lstm_model()
    model.load_weights('lstm_weights.h5')
    result = model.predict(data)
    return result

验证步骤 #

1. 先跑单进程预测确认模型本身没问题
2. 启用spawn模式，在子进程内加载模型
3. 用少量数据测试多进程预测，看是否还会挂起或报错

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