TensorFlow模型恢复耗时过长，求预测场景下的优化方案

Hey there! Let's tackle your slow model restore issue when moving from Colab GPU to local PC, especially since you only need the model for prediction (no further training). First, let's fix small issues in your current code, then dive into optimized solutions.

Quick Fixes to Your Current Code #

Before optimizing, let's clean up your restore code to cut unnecessary overhead:

1. Duplicate Saver Creation: You don't need to initialize tf.train.Saver() before import_meta_graph—the latter returns a valid saver instance automatically.
2. Dropout for Inference: Set keep_prob=1.0 during prediction! Using 0.8 will randomly drop neurons, hurting accuracy and adding unnecessary computation.
3. Unnecessary Nodes: Your restore code includes training-only components (like Y_ placeholder, cross-entropy calculation) which bloat the graph and slow down loading.

Optimized Approaches for Inference-Only Use #

1. Freeze the Graph (Fastest Loading) #

Freezing converts your model's variables into constants and merges the graph structure + weights into a single .pb file. This eliminates the need to re-build the graph and load variables separately—you just load the entire frozen graph directly.

Step 1: Modify Training Code to Save Frozen Graph

# Add this after your graph definition in training code
graph_def = tf.get_default_graph().as_graph_def()

with tf.Session() as sess:
    sess.run(init)
    # ... your existing training loop ...
    
    # Save checkpoint (as you did before)
    save_path = saver.save(sess, "abc/model")
    
    # Freeze the graph (convert variables to constants)
    frozen_graph_def = tf.graph_util.convert_variables_to_constants(
        sess,
        graph_def,
        output_node_names=["Softmax"]  # Replace with your actual output node name (print Y.name during training to confirm)
    )
    
    # Save frozen graph to file
    with open("abc/frozen_model.pb", "wb") as f:
        f.write(frozen_graph_def.SerializeToString())

Step 2: Load Frozen Graph for Inference

import tensorflow as tf

def load_frozen_graph(pb_path):
    with tf.gfile.GFile(pb_path, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name="")
        
        # Fetch input/output tensors (match names from your training graph)
        X = graph.get_tensor_by_name("Placeholder:0")  # Your input placeholder name
        Y = graph.get_tensor_by_name("Softmax:0")  # Your output node name
        # Get dropout keep_prob tensors to set to 1.0
        keep_prob1 = graph.get_tensor_by_name("dropout/keep_prob:0")
        keep_prob2 = graph.get_tensor_by_name("dropout_1/keep_prob:0")
        
        return graph, X, Y, keep_prob1, keep_prob2

# Load the frozen model
graph, X, Y, keep_prob1, keep_prob2 = load_frozen_graph("abc/frozen_model.pb")

# Run inference
with tf.Session(graph=graph) as sess:
    predictions = sess.run(Y, feed_dict={
        X: your_input_data,
        keep_prob1: 1.0,
        keep_prob2: 1.0
    })

2. Use SavedModel Format (TensorFlow's Recommended Deployment Format) #

SavedModel is a standardized, self-contained format that bundles the graph, weights, and metadata. It's designed for easy deployment and loads quickly across environments.

Step 1: Save as SavedModel During Training

with tf.Session() as sess:
    sess.run(init)
    # ... your existing training loop ...
    
    # Save model in SavedModel format
    tf.saved_model.simple_save(
        sess,
        export_dir="abc/saved_model",
        inputs={"input": X},  # Map input placeholder to a friendly name
        outputs={"predictions": Y}  # Map output tensor to a friendly name
    )

Step 2: Load SavedModel for Inference

import tensorflow as tf

with tf.Session() as sess:
    # Load the SavedModel
    tf.saved_model.loader.load(sess, ["serve"], "abc/saved_model")
    
    # Fetch input/output tensors using their friendly names
    X = sess.graph.get_tensor_by_name("input:0")
    Y = sess.graph.get_tensor_by_name("predictions:0")
    keep_prob1 = sess.graph.get_tensor_by_name("dropout/keep_prob:0")
    keep_prob2 = sess.graph.get_tensor_by_name("dropout_1/keep_prob:0")
    
    # Run inference
    predictions = sess.run(Y, feed_dict={
        X: your_input_data,
        keep_prob1: 1.0,
        keep_prob2: 1.0
    })

3. Simplify the Inference Graph (Minimize Overhead) #

If you want to stick with checkpoints, rebuild only the inference-only part of the graph (remove training-specific nodes like Y_, train_step, cross_entropy) to reduce graph size and loading time.

Example Inference-Only Graph

# Only build components needed for prediction
X = tf.placeholder(tf.float32, [None, 56, 56, 1])
L1 = 432
L2 = 72
L3 = 36

# Use the SAME variable names as training to ensure weights load correctly
W1 = tf.Variable(tf.truncated_normal([3136, L1], stddev=0.1), name="W1")
b1 = tf.Variable(tf.zeros([L1]), name="b1")
W2 = tf.Variable(tf.truncated_normal([L1, L2], stddev=0.1), name="W2")
b2 = tf.Variable(tf.zeros([L2]), name="b2")
W3 = tf.Variable(tf.truncated_normal([L2, L3], stddev=0.1), name="W3")
b3 = tf.Variable(tf.zeros([L3]), name="b3")

XX = tf.reshape(X, [-1, 3136])
Y1 = tf.nn.sigmoid(tf.matmul(XX, W1) + b1)
Y1 = tf.nn.dropout(Y1, keep_prob=1.0)  # Disable dropout for inference
Y2 = tf.nn.sigmoid(tf.matmul(Y1, W2) + b2)
Y2 = tf.nn.dropout(Y2, keep_prob=1.0)
Ylogits = tf.matmul(Y2, W3) + b3
Y = tf.nn.softmax(Ylogits)

with tf.Session() as sess:
    saver = tf.train.Saver()
    saver.restore(sess, "abc/model")  # Load weights directly (no meta graph needed if variable names match)
    
    # Run inference
    predictions = sess.run(Y, feed_dict={X: your_input_data})

Why These Methods Are Faster #

• Frozen Graph: Combines graph structure and weights into one file, eliminating variable initialization and graph re-building steps.
• SavedModel: Optimized for deployment, TensorFlow loads it with minimal overhead and handles compatibility automatically.
• Simplified Inference Graph: Removes unnecessary training nodes, reducing the amount of data loaded and processed during restore.

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