TensorFlow：如何将冻结模型（.pb）转换为Saved Model

Absolutely, converting a frozen TensorFlow .pb model to the SavedModel format is fully feasible. SavedModel is TensorFlow's standard serialization format, and the conversion process is straightforward with just a few lines of code.

Here's a step-by-step implementation using TensorFlow 2.x:

Step 1: Prepare the Frozen Model #

First, download and extract your tf_model_pb.zip file to get the frozen .pb file (we’ll assume it’s named frozen_model.pb for this example).

Step 2: Conversion Code #

import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants

# Path to your frozen .pb file
frozen_model_path = "./frozen_model.pb"
# Path to save the converted SavedModel
saved_model_path = "./saved_model"

# Load the frozen graph definition
with tf.io.gfile.GFile(frozen_model_path, "rb") as f:
    graph_def = tf.compat.v1.GraphDef()
    graph_def.ParseFromString(f.read())

# Convert to SavedModel format
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
    # Import the frozen graph into the current session
    tf.import_graph_def(graph_def, name="")
    
    # Replace these with your model's actual input and output tensor names
    # You'll need to know these from your original model setup
    input_tensor = sess.graph.get_tensor_by_name("input:0")
    output_tensor = sess.graph.get_tensor_by_name("output:0")
    
    # Define the inference signature (standard for serving/prediction)
    signature_def = tf.compat.v1.saved_model.signature_def_utils.predict_signature_def(
        inputs={"input": input_tensor},
        outputs={"output": output_tensor}
    )
    
    # Build and save the SavedModel
    builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(saved_model_path)
    builder.add_meta_graph_and_variables(
        sess,
        [tag_constants.SERVING],
        signature_def_map={
            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_def
        }
    )
    builder.save()

print(f"SavedModel successfully saved to {saved_model_path}")

Key Notes: #

• Tensor Names: You must replace "input:0" and "output:0" with the actual names of your model’s input and output tensors. To find these names, run this quick snippet to list all nodes in your frozen graph:

  import tensorflow as tf
  
  frozen_model_path = "./frozen_model.pb"
  with tf.io.gfile.GFile(frozen_model_path, "rb") as f:
      graph_def = tf.compat.v1.GraphDef()
      graph_def.ParseFromString(f.read())
  
  for node in graph_def.node:
      print(node.name)
  

• TensorFlow Compatibility: This code uses TensorFlow 2.x with v1 compatibility modules, which works for most legacy frozen models. If you’re using pure TensorFlow 1.x, you can remove the compat.v1 prefixes where applicable.
• Signature Customization: The example uses a predict_signature_def, which is standard for inference. Adjust the signature type if your model requires classification, regression, or other specialized behavior.

After running this code, you’ll have a SavedModel directory that can be used for local inference, deployment with TensorFlow Serving, or conversion to other formats like TensorFlow Lite.

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