如何修改TensorFlow SavedModel pb文件替换算子以使用自定义op推理

可行性确认 #

你的思路完全可行：只需要把.pb文件中对应的原有算子节点替换为自定义Op节点，再在启动TensorFlow Serving时加载自定义Op的动态库即可，不需要改动训练逻辑。

具体操作步骤 #

1. 提取原有模型的节点特征 #

首先需要确认你要替换的目标算子的标识特征（Op类型、属性、输入输出结构等），避免误替换其他节点，可以用以下代码打印模型节点信息：

import tensorflow as tf
from tensorflow.core.framework import graph_pb2

def load_pb(pb_path):
    graph_def = graph_pb2.GraphDef()
    with open(pb_path, "rb") as f:
        graph_def.ParseFromString(f.read())
    return graph_def

# 加载原模型并打印节点信息
origin_graph = load_pb("你的原模型路径.pb")
for node in origin_graph.node[:100]:
    print(f"节点名: {node.name}, Op类型: {node.op}, 持有属性: {list(node.attr.keys())}")

2. 实现算子替换逻辑 #

以下是单类型算子替换的参考代码，如果是多类型或者子图替换，调整匹配逻辑即可：

def replace_ops(graph_def, origin_op_type, custom_op_type, extra_attrs=None):
    new_graph = graph_pb2.GraphDef()
    extra_attrs = extra_attrs or {}
    for node in graph_def.node:
        new_node = new_graph.node.add()
        # 匹配到目标算子则替换Op类型
        if node.op == origin_op_type:
            new_node.CopyFrom(node)
            new_node.op = custom_op_type
            # 写入自定义Op需要的额外属性
            for k, v in extra_attrs.items():
                if isinstance(v, int):
                    new_node.attr[k].i = v
                elif isinstance(v, float):
                    new_node.attr[k].f = v
                elif isinstance(v, str):
                    new_node.attr[k].s = v.encode()
        # 非目标算子直接拷贝
        else:
            new_node.CopyFrom(node)
    new_graph.versions.CopyFrom(graph_def.versions)
    return new_graph

# 示例：替换所有Conv2D算子为自定义FastConv2D，传入自定义属性
modified_graph = replace_ops(
    graph_def=origin_graph,
    origin_op_type="Conv2D",
    custom_op_type="FastConv2D",
    extra_attrs={"enable_fast_mode": 1}
)

# 保存替换后的模型
with open("替换后模型.pb", "wb") as f:
    f.write(modified_graph.SerializeToString())

3. 本地验证替换效果 #

部署到TensorFlow Serving之前先本地验证推理一致性，避免结构错误：

with tf.Graph().as_default():
    tf.import_graph_def(modified_graph, name="")
    with tf.Session() as sess:
        input_tensor = sess.graph.get_tensor_by_name("你的输入节点名:0")
        output_tensor = sess.graph.get_tensor_by_name("你的输出节点名:0")
        # 构造随机测试输入，和原模型的推理结果做对比
        test_input = tf.random.normal(shape=你的输入维度).eval()
        custom_output = sess.run(output_tensor, feed_dict={input_tensor: test_input})
        # 确认误差在可接受范围内即可

4. TensorFlow Serving加载配置 #

启动TFServing时加上自定义Op动态库的加载参数即可，运行时会自动匹配.pb中的自定义Op类型调用对应实现：

tensorflow_model_server \
  --model_name=你的模型名 \
  --model_base_path=你的模型路径 \
  --custom_op_paths=你的自定义Op编译后的.so文件路径

注意事项 #

• 自定义Op的输入输出数量、数据类型必须和原有算子完全一致，否则会出现图结构不兼容报错
• 如果是替换多节点组成的子图，只需要删掉原有子图的所有节点，插入自定义Op节点后重新连接输入输出关系即可，逻辑和单节点替换一致
• 自定义Op不需要兼容训练逻辑，只要推理阶段的前向逻辑正确即可正常运行

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