Ray.io分布式TensorFlow推理时Graph对象序列化失败问题求助

问题核心原因 #

序列化错误的本质是TensorFlow的Graph对象内部包含_thread.RLock这类无法被pickle序列化的线程同步对象，而Ray在分发远程任务时需要序列化所有传入参数，导致无法传递Graph对象。

可行解决方案 #

方案1：在Ray远程函数内部加载模型/Graph #

不要将Graph作为参数传递给远程函数，而是在每个远程任务内部完成模型和Graph的加载。这样既避免了序列化问题，还能让每个Ray worker独立持有模型实例，适配CPU/GPU分布式场景。

修改后的代码示例：

N_CHANNELS = 3
import ray
import time
import numpy as np
from PIL import Image
import tensorflow as tf

def load_image_into_numpy_array(image):
    """
    Converts a PIL image into a numpy array (height x width x channels).
    :param image: PIL image
    :return: numpy array
    """
    (width, height) = image.size
    return np.array(image.getdata()) \
        .reshape((height, width, N_CHANNELS)).astype(np.uint8)

# 新增模型加载函数，替换为你的实际模型加载逻辑
def load_tf_model():
    graph = tf.Graph()
    with graph.as_default():
        # 这里添加你的模型加载逻辑，比如从pb文件导入
        # tf.compat.v1.train.import_meta_graph("your_model_path.pb")
        pass
    return graph

@ray.remote(num_gpus=1)  # 使用GPU时指定资源
def run_inference(image_np):
    """
    Runs the inference on the given image.
    :param image_np: numpy image
    :return: dictionary with detected classes 
             and their corresponding scores and boxes
    """
    # 在远程函数内部加载Graph
    graph = load_tf_model()
    output_tensor_dict = {
        DETECTION_BOXES_KEY: DETECTION_BOXES_KEY + TENSOR_SUFFIX,
        DETECTION_SCORES_KEY: DETECTION_SCORES_KEY + TENSOR_SUFFIX,
        DETECTION_CLASSES_KEY: DETECTION_CLASSES_KEY + TENSOR_SUFFIX
    }

    with graph.as_default():
        with tf.compat.v1.Session() as sess:
            input_tensor = tf.compat.v1.get_default_graph()\
                .get_tensor_by_name(IMAGE_TENSOR_KEY + TENSOR_SUFFIX)
            image_np_expanded = np.expand_dims(image_np, axis=0)
            input_tensor_dict = {input_tensor: image_np_expanded}
            output_dict = sess.run(output_tensor_dict,
                                   feed_dict=input_tensor_dict)

            return {
                DETECTION_BOXES_KEY: 
                    output_dict[DETECTION_BOXES_KEY][0],
                DETECTION_SCORES_KEY: 
                    output_dict[DETECTION_SCORES_KEY][0],
                DETECTION_CLASSES_KEY: 
                    output_dict[DETECTION_CLASSES_KEY][0].astype(np.int64)
            }

# =====================
# Running the inference
# =====================
tic = time.perf_counter()
print("* Stage 9: Running the inference")
IMAGE_NP_KEY = 'image_np'
RESULTS_KEY = 'results'

file_result_dict = {}
ray.init()

# 批量提交任务，不再传递graph参数
futures = [
    run_inference.remote(load_image_into_numpy_array(Image.open(filename))) 
    for filename in TEST_IMAGES
]
# 等待所有任务完成
output_dicts = ray.get(futures)

# 处理结果
for filename, output_dict in zip(TEST_IMAGES, output_dicts):
    results = process_output(
        output_dict[DETECTION_CLASSES_KEY],
        output_dict[DETECTION_SCORES_KEY],
        output_dict[DETECTION_BOXES_KEY],
        category_index
    )
    file_result_dict[filename] = { 
        IMAGE_NP_KEY: load_image_into_numpy_array(Image.open(filename)), 
        RESULTS_KEY: results 
    }

toc = time.perf_counter()
print("Stage 9 completed in", round(toc - tic, 2), "seconds")
ray.shutdown()

方案2：使用Ray Actor复用模型实例 #

如果模型加载成本高（比如大模型），可以用Ray Actor在进程初始化时加载一次模型，之后处理多个任务，避免重复加载，提升运行效率。

代码示例：

import ray
import tensorflow as tf
import numpy as np
from PIL import Image

@ray.remote(num_gpus=1)
class InferenceActor:
    def __init__(self):
        # Actor初始化时加载模型和Graph，仅执行一次
        self.graph = load_tf_model()
        self.sess = tf.compat.v1.Session(graph=self.graph)
        with self.graph.as_default():
            self.input_tensor = tf.compat.v1.get_default_graph()\
                .get_tensor_by_name(IMAGE_TENSOR_KEY + TENSOR_SUFFIX)
        self.output_tensor_dict = {
            DETECTION_BOXES_KEY: DETECTION_BOXES_KEY + TENSOR_SUFFIX,
            DETECTION_SCORES_KEY: DETECTION_SCORES_KEY + TENSOR_SUFFIX,
            DETECTION_CLASSES_KEY: DETECTION_CLASSES_KEY + TENSOR_SUFFIX
        }

    def run_inference(self, image_np):
        image_np_expanded = np.expand_dims(image_np, axis=0)
        input_tensor_dict = {self.input_tensor: image_np_expanded}
        output_dict = self.sess.run(self.output_tensor_dict,
                                   feed_dict=input_tensor_dict)
        return {
            DETECTION_BOXES_KEY: 
                output_dict[DETECTION_BOXES_KEY][0],
            DETECTION_SCORES_KEY: 
                output_dict[DETECTION_SCORES_KEY][0],
            DETECTION_CLASSES_KEY: 
                output_dict[DETECTION_CLASSES_KEY][0].astype(np.int64)
        }

# 运行逻辑
ray.init()
# 创建Actor实例（每个Actor占用一个GPU资源）
actor = InferenceActor.remote()

# 提交任务
futures = [
    actor.run_inference.remote(load_image_into_numpy_array(Image.open(filename))) 
    for filename in TEST_IMAGES
]
output_dicts = ray.get(futures)

# 后续结果处理逻辑与方案1一致
ray.shutdown()

方案3：迁移到TensorFlow 2.x的SavedModel格式 #

既然使用的是TF2.9.1，建议放弃compat.v1的Graph和Session，改用TF2原生的SavedModel格式，这种格式的模型可直接序列化，更适配Ray的分布式场景。

示例代码（简化版）：

import ray
import tensorflow as tf
import numpy as np
from PIL import Image

# 训练阶段导出模型为SavedModel格式
# tf.saved_model.save(model, "./saved_model")

@ray.remote(num_gpus=1)
def run_inference(image_np):
    # 加载SavedModel
    model = tf.saved_model.load("./saved_model")
    infer = model.signatures["serving_default"]
    image_np_expanded = np.expand_dims(image_np, axis=0)
    output_dict = infer(tf.convert_to_tensor(image_np_expanded))
    # 转换结果格式
    return {
        DETECTION_BOXES_KEY: output_dict[DETECTION_BOXES_KEY].numpy()[0],
        DETECTION_SCORES_KEY: output_dict[DETECTION_SCORES_KEY].numpy()[0],
        DETECTION_CLASSES_KEY: output_dict[DETECTION_CLASSES_KEY].numpy()[0].astype(np.int64)
    }

方案选择建议 #

• 模型较小、加载速度快：优先用方案1，实现简单直接。
• 模型体积大、加载成本高：优先用方案2，通过Actor复用模型，减少重复加载开销。
• 长期维护的项目：优先用方案3，贴合TF2生态，避免compat.v1的兼容问题。

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