Keras多进程模型预测卡住问题及相关技术疑问

• 有一个MNIST Keras模型，用于执行预测并计算损失值，在多CPU服务器上尝试用多进程提速，但进程永远无法结束，单进程运行正常。
• 尝试过在每个进程内加载模型，或使用全局模型，均无效，predict函数中的打印语句从未执行。
• 代码示例如下：

from multiprocessing import Process
import tensorflow as tf

#make a prediction on a training sample
def predict(idx, return_dict):
  x = tf.convert_to_tensor(np.expand_dims(x_train[idx],axis=0))

  local_model=tf.keras.models.load_model('model.h5')
  y=local_model(x)
  print('this never gets printed')
  y_expanded=np.expand_dims(y_train[train_idx],axis=0)
  loss=tf.keras.losses.CategoricalCrossentropy(y_expanded,y)
  return_dict[i]=loss

manager = multiprocessing.Manager()
return_dict = manager.dict()
jobs = []

for i in range(10):
    p = Process(target=predict, args=(i, return_dict))
    jobs.append(p)
    p.start()
    
for proc in jobs:
    proc.join()

print(return_dict.values())

• 疑问：
  1. 如何解决模型导致的多进程阻塞问题
  2. 是否可以让所有进程共用同一个X_train

1. 解决模型多进程阻塞问题 #

TensorFlow多进程阻塞的核心原因是父进程的TensorFlow会话/计算图被子进程继承，导致资源冲突。结合代码中的变量错误，可按以下步骤修复：

关键修复点 #

• 在子进程内重置TensorFlow上下文：每个进程加载模型前，执行tf.keras.backend.clear_session()，确保进程拥有独立的TF运行环境，避免父进程资源干扰。
• 修正代码变量错误：代码中train_idx应为idx，return_dict[i]应为return_dict[idx]，这些隐性错误会导致进程异常但无报错输出。
• 规范多进程启动方式：将主逻辑放在if __name__ == '__main__':块内，这是Python多进程的强制规范，避免子进程重复执行主代码引发异常。
• 改用Pool管理进程：Pool会自动处理进程初始化逻辑，减少手动管理的出错概率，还可设置maxtasksperchild防止内存泄漏。

修正后的代码示例 #

import numpy as np
import tensorflow as tf
from multiprocessing import Manager, Pool

def predict(idx, return_dict, x_train, y_train):
    # 重置TF上下文，创建独立环境
    tf.keras.backend.clear_session()
    # 转换输入格式
    x = tf.convert_to_tensor(np.expand_dims(x_train[idx], axis=0))
    # 加载模型
    local_model = tf.keras.models.load_model('model.h5')
    y = local_model(x)
    print(f'完成预测，索引：{idx}')
    # 修正变量名错误
    y_expanded = np.expand_dims(y_train[idx], axis=0)
    # 正确调用损失函数：先实例化类，再计算损失
    loss_fn = tf.keras.losses.CategoricalCrossentropy()
    loss = loss_fn(y_expanded, y)
    # 转成numpy值存入字典，避免TF张量序列化问题
    return_dict[idx] = loss.numpy()

if __name__ == '__main__':
    # 主进程加载数据集（示例用MNIST官方数据集）
    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
    x_train = x_train.astype('float32') / 255.0
    y_train = tf.keras.utils.to_categorical(y_train, 10)
    
    manager = Manager()
    return_dict = manager.dict()
    
    # 使用Pool管理进程，进程数可设为CPU核心数
    with Pool(processes=4) as pool:
        for i in range(10):
            pool.apply_async(predict, args=(i, return_dict, x_train, y_train))
        pool.close()
        pool.join()
    
    print(return_dict.values())

2. 让所有进程共用X_train #

可以通过两种方式实现，避免每个进程重复加载数据：

• Unix系统：利用写时复制机制：在主进程加载X_train，子进程通过fork机制自动共享数据（写时复制，子进程不修改数据则不会复制内存），无需额外操作，节省内存开销。
• 跨平台：使用共享内存：对于Windows系统（采用spawn机制，不会自动共享父进程数据），可使用multiprocessing.Array或numpy共享内存数组，将X_train转换成共享内存对象，所有进程直接读取。

跨平台共享内存示例（简化版） #

import numpy as np
import tensorflow as tf
from multiprocessing import Manager, Pool, Array

def shared_array_to_numpy(shared_arr, shape):
    # 将共享内存数组转换为numpy数组
    return np.frombuffer(shared_arr.get_obj(), dtype=np.float32).reshape(shape)

if __name__ == '__main__':
    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
    x_train = x_train.astype('float32') / 255.0
    y_train = tf.keras.utils.to_categorical(y_train, 10)
    
    # 将x_train转为共享内存数组
    x_shape = x_train.shape
    shared_x = Array('f', x_train.flatten(), lock=False)
    
    manager = Manager()
    return_dict = manager.dict()
    
    def predict_shared(idx, return_dict):
        tf.keras.backend.clear_session()
        # 从共享内存读取x_train
        x_np = shared_array_to_numpy(shared_x, x_shape)
        x = tf.convert_to_tensor(np.expand_dims(x_np[idx], axis=0))
        local_model = tf.keras.models.load_model('model.h5')
        y = local_model(x)
        print(f'完成预测，索引：{idx}')
        y_expanded = np.expand_dims(y_train[idx], axis=0)
        loss_fn = tf.keras.losses.CategoricalCrossentropy()
        loss = loss_fn(y_expanded, y)
        return_dict[idx] = loss.numpy()
    
    with Pool(processes=4) as pool:
        for i in range(10):
            pool.apply_async(predict_shared, args=(i, return_dict))
        pool.close()
        pool.join()
    
    print(return_dict.values())

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