TensorFlow新手求助：从预训练模型生成冻结图并转.tflite

Hey there! I totally get wanting to walk through these steps manually even though the model comes with pre-built files—nothing beats hands-on practice to solidify your TensorFlow understanding. Let’s break this down step by step, with code examples and key notes to avoid pitfalls.

Step 1: Generate a Frozen Graph from the CKPT Files #

A frozen graph combines your model’s structure (from the .meta file) and trained weights (from .ckpt.data/.ckpt.index) into a single .pb file, with all variables converted to constants. Here’s how to do it:

First, create a Python script (e.g., freeze_graph.py):

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

# Configure paths to your model files
CKPT_PATH = "./mobilenet_v1_1.0_224.ckpt"
META_PATH = "./mobilenet_v1_1.0_224.ckpt.meta"
OUTPUT_FROZEN_PB = "./my_custom_frozen_graph.pb"

# Critical: Define your model's output node name(s)
# For Mobilenet V1, this is typically 'MobilenetV1/Predictions/Softmax'
OUTPUT_NODE_NAMES = "MobilenetV1/Predictions/Softmax"

with tf.compat.v1.Session() as sess:
    # Load the model graph structure from the .meta file
    saver = tf.compat.v1.train.import_meta_graph(META_PATH)
    # Restore the trained weights from the CKPT files
    saver.restore(sess, CKPT_PATH)
    
    # Get the current graph definition
    graph_def = tf.compat.v1.get_default_graph().as_graph_def()
    
    # Freeze the graph: convert variables to constants, keep only output-related nodes
    frozen_graph_def = graph_util.convert_variables_to_constants(
        sess, graph_def, OUTPUT_NODE_NAMES.split(",")
    )
    
    # Save the frozen graph to disk
    with tf.io.gfile.GFile(OUTPUT_FROZEN_PB, "wb") as f:
        f.write(frozen_graph_def.SerializeToString())
    
    print(f"Success! Frozen graph saved to {OUTPUT_FROZEN_PB}")

Key Notes for Step 1: #

• Finding the output node name: If you’re unsure, run [n.name for n in tf.compat.v1.get_default_graph().as_graph_def().node] after loading the graph to print all node names, or check the mobilenet_v1_1.0_224_eval.pbtxt file for the prediction node.
• TF 2.x Compatibility: Since Mobilenet V1 is a TF1.x model, we use tf.compat.v1 to maintain compatibility.

Step 2: Convert the Frozen Graph to TFLite Format #

Once you have your frozen .pb file, you can convert it to TFLite using either the Python API or command line. Both methods work—pick whichever you prefer for learning.

Option A: Python API (Great for scripting and customization) #

Create another script (e.g., convert_to_tflite.py):

import tensorflow as tf

# Configure paths and node names
FROZEN_PB_PATH = "./my_custom_frozen_graph.pb"
OUTPUT_TFLITE = "./my_custom_mobilenet.tflite"
INPUT_NODE_NAME = "input"  # Mobilenet V1's standard input node
OUTPUT_NODE_NAME = "MobilenetV1/Predictions/Softmax"

# Initialize the converter with the frozen graph
converter = tf.lite.TFLiteConverter.from_frozen_graph(
    graph_def_file=FROZEN_PB_PATH,
    input_arrays=[INPUT_NODE_NAME],
    input_shapes={INPUT_NODE_NAME: [1, 224, 224, 3]},  # Shape: [batch_size, height, width, channels]
    output_arrays=[OUTPUT_NODE_NAME]
)

# Optional: Enable post-training quantization to reduce model size (great for edge devices)
# converter.optimizations = [tf.lite.Optimize.DEFAULT]

# Perform the conversion
tflite_model = converter.convert()

# Save the TFLite model
with open(OUTPUT_TFLITE, "wb") as f:
    f.write(tflite_model)

print(f"Success! TFLite model saved to {OUTPUT_TFLITE}")

Option B: Command Line (Quick for one-off conversions) #

Run this command in your terminal:

tflite_convert \
  --graph_def_file=./my_custom_frozen_graph.pb \
  --output_file=./my_custom_mobilenet.tflite \
  --input_arrays=input \
  --input_shapes=1,224,224,3 \
  --output_arrays=MobilenetV1/Predictions/Softmax

Key Notes for Step 2: #

• Input shape: The [1,224,224,3] shape corresponds to a single 224x224 RGB image. Adjust the batch size (first dimension) if needed.
• Quantization: The optional optimization line reduces model size without significant accuracy loss—definitely worth experimenting with once you’ve mastered the basic conversion.

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