如何拆分大体积TensorFlow SavedModel为两部分并分存部署到CPU与GPU？

嘿，这个场景我之前在做推荐系统模型优化的时候正好碰到过！拆分SavedModel来分别部署CPU和GPU负载，确实是解决这类内存密集+计算密集混合模型的最优方案之一。我来给你一步步拆解怎么做，包括具体的代码示例和踩过坑后的最佳实践。

你的模型分界点非常明确——Gather、SparseSegmentSum、SparseSegmentMean这些操作就是上游稀疏嵌入模块（CPU友好）和下游DNN模块（GPU友好）的接口。拆分的关键是：

• 上游模块：接受原始稀疏输入，输出经过聚合后的稠密嵌入向量（也就是分界操作的结果）
• 下游模块：以上游输出的嵌入向量为输入，执行DNN计算并输出最终预测结果

这样拆分后，体积庞大的稀疏嵌入矩阵只存在于上游模型，部署在内存充足的CPU机器上；下游模型体积小，专门放GPU上跑计算密集的深度网络。

我用TensorFlow 2.x来写示例，这也是当前的主流版本，如果你还在使用TF1.x，思路类似但API会有差异。

1. 加载原始SavedModel #

首先把你的大模型加载进来，拿到它的推理签名：

import tensorflow as tf

# 替换成你的原始模型路径
original_model_path = "/path/to/your/large_savedmodel"
original_model = tf.saved_model.load(original_model_path)
# 假设你的模型用的是默认的serving签名，根据实际情况调整
infer_signature = original_model.signatures["serving_default"]

2. 导出上游CPU模型 #

上游模型的核心是处理稀疏输入，生成聚合后的嵌入向量。这里有两种实现方式：

方式一：直接从原始计算图提取子图（快速但依赖内部命名） #

如果能找到分界操作的输出张量名，可以直接提取子图：

# 获取原始模型的 concrete function，方便访问计算图
concrete_fn = infer_signature.get_concrete_function()
graph = concrete_fn.graph

# 替换成你实际的分界操作输出张量名，比如通过tensorboard查看计算图找到
aggregated_embeddings_tensor = graph.get_tensor_by_name("aggregated_embeddings:0")

# 定义上游模型的推理函数，输入和原始模型一致，输出为聚合后的嵌入
@tf.function(input_signature=infer_signature.input_signature)
def upstream_infer(inputs):
    # 执行原始模型前向传播到分界点
    return {"aggregated_embeddings": aggregated_embeddings_tensor}

# 导出上游SavedModel
upstream_model_path = "/path/to/upstream_cpu_model"
tf.saved_model.save(
    upstream_infer,
    upstream_model_path,
    signatures={"serving_default": upstream_infer}
)

方式二：重新构建上游逻辑（更可靠，推荐） #

直接依赖内部张量名容易因为模型导出时的命名变化出错，更稳妥的方式是重新实现上游的稀疏嵌入+聚合逻辑，然后加载原始权重：

# 假设你的上游结构是：稀疏嵌入Lookup → Gather → SparseSegmentSum
# 1. 定义上游模型结构
class UpstreamModel(tf.keras.Model):
    def __init__(self, vocab_size, embedding_dim):
        super().__init__()
        # 稀疏嵌入矩阵，和原始模型的嵌入层一致
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        # 其他聚合操作可以在call里实现

    def call(self, inputs):
        # 假设输入是稀疏张量的三个组件：indices, values, dense_shape
        sparse_indices = inputs["sparse_indices"]
        sparse_values = inputs["sparse_values"]
        dense_shape = inputs["dense_shape"]
        
        # 构建稀疏张量
        sparse_tensor = tf.sparse.SparseTensor(sparse_indices, sparse_values, dense_shape)
        # 嵌入Lookup
        embeddings = self.embedding(sparse_tensor.indices[:, 1])  # 根据你的输入结构调整
        # 执行SparseSegmentSum聚合
        aggregated_embeddings = tf.sparse.segment_sum(
            data=embeddings,
            segment_ids=sparse_tensor.indices[:, 0],
            num_segments=dense_shape[0]
        )
        return {"aggregated_embeddings": aggregated_embeddings}

# 2. 初始化上游模型，替换成你的实际参数
upstream_model = UpstreamModel(vocab_size=1000000, embedding_dim=128)
# 3. 从原始模型加载嵌入层权重
upstream_model.embedding.set_weights([original_model.embedding.weights[0].numpy()])

# 4. 导出SavedModel，指定输入签名（和原始模型一致）
input_signature = [
    tf.TensorSpec(shape=[None, 2], dtype=tf.int64, name="sparse_indices"),
    tf.TensorSpec(shape=[None], dtype=tf.float32, name="sparse_values"),
    tf.TensorSpec(shape=[2], dtype=tf.int64, name="dense_shape")
]
# 或者直接用原始模型的输入签名
upstream_model.save(
    upstream_model_path,
    save_format="tf",
    signatures={"serving_default": tf.function(input_signature=input_signature)(upstream_model)}
)

3. 导出下游GPU模型 #

下游模型的输入是上游的嵌入向量，输出是最终预测。同样推荐重新构建结构并加载权重：

# 1. 定义下游DNN结构，和原始模型一致
class DownstreamModel(tf.keras.Model):
    def __init__(self, embedding_dim, num_classes):
        super().__init__()
        self.dense1 = tf.keras.layers.Dense(512, activation="relu")
        self.dropout = tf.keras.layers.Dropout(0.5)
        self.dense2 = tf.keras.layers.Dense(256, activation="relu")
        self.output_layer = tf.keras.layers.Dense(num_classes, activation="softmax")

    def call(self, inputs):
        x = self.dense1(inputs["aggregated_embeddings"])
        x = self.dropout(x)
        x = self.dense2(x)
        return {"predictions": self.output_layer(x)}

# 2. 初始化下游模型
downstream_model = DownstreamModel(embedding_dim=128, num_classes=10)
# 3. 从原始模型加载下游层的权重
# 假设原始模型的下游层变量名是dense1, dropout, dense2, output_layer
downstream_model.dense1.set_weights([original_model.dense1.weights[0].numpy(), original_model.dense1.weights[1].numpy()])
downstream_model.dense2.set_weights([original_model.dense2.weights[0].numpy(), original_model.dense2.weights[1].numpy()])
downstream_model.output_layer.set_weights([original_model.output_layer.weights[0].numpy(), original_model.output_layer.weights[1].numpy()])

# 4. 导出下游SavedModel
downstream_model_path = "/path/to/downstream_gpu_model"
input_signature = [tf.TensorSpec(shape=[None, 128], dtype=tf.float32, name="aggregated_embeddings")]
downstream_model.save(
    downstream_model_path,
    save_format="tf",
    signatures={"serving_default": tf.function(input_signature=input_signature)(downstream_model)}
)

4. 验证拆分后的模型 #

确保两个模型串联后的输出和原始模型一致：

# 用CPU跑上游
with tf.device("/CPU:0"):
    upstream = tf.saved_model.load(upstream_model_path)
    upstream_infer = upstream.signatures["serving_default"]
    # 模拟稀疏输入
    test_input = {
        "sparse_indices": tf.convert_to_tensor([[0, 1], [0, 3], [1, 2]], dtype=tf.int64),
        "sparse_values": tf.convert_to_tensor([1.0, 2.0, 3.0], dtype=tf.float32),
        "dense_shape": tf.convert_to_tensor([2, 4], dtype=tf.int64)
    }
    embeddings = upstream_infer(**test_input)["aggregated_embeddings"]

# 用GPU跑下游
with tf.device("/GPU:0"):
    downstream = tf.saved_model.load(downstream_model_path)
    downstream_infer = downstream.signatures["serving_default"]
    split_predictions = downstream_infer(aggregated_embeddings=embeddings)["predictions"]

# 和原始模型输出对比
original_predictions = infer_signature(**test_input)["predictions"]
assert tf.reduce_all(tf.abs(split_predictions - original_predictions) < 1e-6), "输出不一致，请检查拆分逻辑！"

• 优先重新构建子图：直接提取计算图的方式虽然快，但依赖内部张量命名，模型迭代后很容易失效。重新构建结构+加载原始权重的方式更可控，也便于后续维护和修改。
• 设备约束与部署：
  • 部署上游模型时，强制绑定CPU（比如TensorFlow Serving配置中指定--enable_cpu_placement=true），避免它占用GPU内存。
  • 下游模型部署在GPU机器上，开启XLA优化（tf.config.optimizer.set_jit(True)）进一步提升性能。
• 架构选型：
  • 用TensorFlow Serving分别部署两个模型，通过API网关（比如Nginx）串联，或者用gRPC直接在服务间通信，减少数据传输开销。
  • 如果是Kubernetes环境，可以给上游服务分配CPU资源，下游服务分配GPU资源，实现弹性扩缩容。
• 性能调优：
  • 上游模型：调整批量大小以匹配CPU内存，开启稀疏张量的优化（比如TF的tf.sparse模块的高效实现）。
  • 下游模型：根据GPU显存调整批量大小，开启混合精度训练/推理（tf.keras.mixed_precision.set_global_policy('mixed_float16')）。
• 版本同步：拆分后的两个模型要保持版本一致，避免因为权重不匹配导致输出错误，可以用统一的版本号命名模型目录。

• 避免变量共享：如果原始模型上下游有共享变量，拆分时要确保变量只存在于一个模型中，或者通过权重同步保持一致。
• 稀疏转稠密：上游输出必须是稠密张量，GPU处理稀疏张量效率极低，所以一定要在CPU端完成稀疏聚合的操作。
• TF版本兼容：如果用TF1.x，需要用SavedModelBuilder来导出模型，并且要处理会话和计算图的问题，推荐升级到TF2.x简化流程。

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