使用自定义算子冻结TensorFlow计算图时遇加载问题求助

这个问题我之前也碰到过！核心原因是TensorFlow导出计算图的时候，并不会把自定义算子的实现打包到导出的.pb文件或者SavedModel里——它只保存了算子的名称和输入输出信息，所以重新加载图的时候，必须先让TensorFlow知道这个自定义算子的存在，也就是提前加载对应的.so库。

下面给你具体的解决步骤和代码示例：

1. 先完善你的导出代码 #

首先确保导出图的流程是完整的，并且先加载自定义算子库再构建计算图：

import tensorflow as tf
from tensorflow.python.framework import graph_util
from tensorflow.python.framework import graph_io

# 关键第一步：先加载自定义算子库，再构建图
custom_op_lib = tf.load_op_library('/path/to/your/custom_op.so')

with tf.device('/gpu:0'):
    with tf.Session() as sess:
        # 构建包含自定义算子的计算图
        input_tensor = tf.placeholder(tf.float32, shape=[None, 224, 224, 3], name='input')
        # 这里替换成你的自定义算子调用
        output_tensor = custom_op_lib.your_custom_op(input_tensor, name='output')

        # 初始化所有变量（如果你的图里有变量的话）
        sess.run(tf.global_variables_initializer())

        # 导出冻结图（把变量转成常量）
        frozen_graph_def = graph_util.convert_variables_to_constants(
            sess,
            sess.graph_def,
            output_node_names=['output']  # 务必指定你的输出节点名称
        )
        # 保存到本地
        graph_io.write_graph(frozen_graph_def, './exported_graph', 'custom_op_graph.pb', as_text=False)

2. 重新加载图的正确姿势 #

加载图的时候，必须先加载自定义算子库，再导入图，否则TensorFlow会识别不了自定义算子：

import tensorflow as tf

# 关键：先加载自定义算子库，顺序不能错！
custom_op_lib = tf.load_op_library('/path/to/your/custom_op.so')

with tf.Session() as sess:
    # 读取冻结图文件
    with tf.gfile.GFile('./exported_graph/custom_op_graph.pb', 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())

    # 导入图到当前会话
    input_tensor, output_tensor = tf.import_graph_def(
        graph_def,
        return_elements=['input:0', 'output:0']  # 对应导出时的节点名
    )

    # 测试运行
    test_input = tf.random_normal([1, 224, 224, 3]).eval()
    result = sess.run(output_tensor, feed_dict={input_tensor: test_input})
    print("自定义算子运行结果：", result)

3. 常见坑点提醒 #

• 加载顺序不能乱：一定要先加载.so库，再加载计算图，否则TensorFlow遇到未知算子直接报错。
• 版本兼容：导出图时用的TensorFlow版本，和加载图时的版本要一致，否则可能出现ABI不兼容的问题（比如算子的底层调用接口变了）。
• 路径问题：尽量用绝对路径加载.so文件，避免相对路径导致找不到文件的情况。
• 如果用SavedModel格式：导出和加载时同样要遵循先加载算子库的规则，这里给个快速示例：

  # 导出SavedModel
  builder = tf.saved_model.builder.SavedModelBuilder('./saved_model')
  builder.add_meta_graph_and_variables(
      sess,
      [tf.saved_model.tag_constants.SERVING],
      signature_def_map={
          'serving_default': tf.saved_model.signature_def_utils.predict_signature_def(
              inputs={'input': input_tensor},
              outputs={'output': output_tensor}
          )
      }
  )
  builder.save()
  
  # 加载SavedModel
  custom_op_lib = tf.load_op_library('/path/to/your/custom_op.so')
  with tf.Session(graph=tf.Graph()) as sess:
      tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], './saved_model')
      input_tensor = sess.graph.get_tensor_by_name('input:0')
      output_tensor = sess.graph.get_tensor_by_name('output:0')
      # 运行测试...
  

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