如何使用tf.layers在预训练模型输出层前插入新层（含CNN场景）

Alright, let's tackle this step by step—first I'll cover the general method for inserting a new layer before the output layer in a pre-trained model with tf.layers, then dive into the exact implementation tailored to your CNN setup.

General Approach #

When working with pre-trained models in TensorFlow (using tf.layers), the core steps to insert a new layer while freezing existing weights are:

• Reconstruct the original model structure: You need to define the exact same layers as your pre-trained model (matching layer names perfectly) so TensorFlow can map the saved weights correctly.
• Freeze existing layers: Set the trainable attribute of all original layers to False so their pre-trained weights don't update during subsequent training.
• Insert the new layer: Connect the new layer to the output of the layer right before your target insertion point (FC1 in your case), then rewire the original output layer to take input from this new layer.
• Load pre-trained weights: Use a Saver that only loads variables from the original model, then initialize the new layer's weights separately to avoid loading errors.

Exact Implementation for Your CNN Scenario #

You mentioned your model has 2 conv layers, an FC1 layer, and a 2-node output layer, saved via saver.save(sess, path). Here's how to modify it properly:

Step 1: Reconstruct the Original Model #

First, define the original layer structure with matching names (this is critical for loading weights correctly):

import tensorflow as tf

def build_original_model(inputs):
    # Conv Layer 1
    conv1 = tf.layers.conv2d(inputs, 32, 3, activation='relu', name='conv1')
    conv1_pool = tf.layers.max_pooling2d(conv1, 2, 2, name='conv1_pool')
    # Conv Layer 2
    conv2 = tf.layers.conv2d(conv1_pool, 64, 3, activation='relu', name='conv2')
    conv2_pool = tf.layers.max_pooling2d(conv2, 2, 2, name='conv2_pool')
    # Flatten for fully connected layers
    flatten = tf.layers.flatten(conv2_pool, name='flatten')
    # FC1 Layer
    fc1 = tf.layers.dense(flatten, 128, activation='relu', name='fc1')
    # Original Output Layer
    original_output = tf.layers.dense(fc1, 2, name='original_output')
    return fc1, original_output

Step 2: Build the Graph & Freeze Original Layers #

Now construct the graph, then freeze all original layers to lock their pre-trained weights:

# Define input placeholder (adjust shape to match your dataset)
inputs = tf.placeholder(tf.float32, shape=[None, 28, 28, 1])
fc1, original_output = build_original_model(inputs)

# Freeze all original layers: set their trainable attribute to False
# Option 1: Target specific variables by exact name
for var_name in ['conv1/kernel:0', 'conv1/bias:0',
                 'conv2/kernel:0', 'conv2/bias:0',
                 'fc1/kernel:0', 'fc1/bias:0',
                 'original_output/kernel:0', 'original_output/bias:0']:
    tf.get_default_graph().get_tensor_by_name(var_name).trainable = False

# Option 2: More concise - filter variables by prefix (easier if you have many layers)
# for var in tf.trainable_variables():
#     if var.name.startswith(('conv1/', 'conv2/', 'fc1/', 'original_output/')):
#         var.trainable = False

Step 3: Insert the New FC2 Layer & Rewire the Output #

Add your new FC2 layer, then reconnect the original output layer to use FC2's output instead of FC1's:

# Insert new FC2 layer (trainable by default, so it will update during training)
fc2 = tf.layers.dense(fc1, 64, activation='relu', name='fc2')

# Rewire the original output layer to take input from FC2
# We'll reuse the original output layer's pre-trained weights to avoid retraining it
output_kernel = tf.get_default_graph().get_tensor_by_name('original_output/kernel:0')
output_bias = tf.get_default_graph().get_tensor_by_name('original_output/bias:0')
new_output = tf.matmul(fc2, output_kernel) + output_bias

Step 4: Load Pre-trained Weights & Initialize New Layer #

Finally, load the saved weights (only for original layers) and initialize the new FC2 layer:

# Initialize only the new FC2 layer's variables (avoid overwriting pre-trained weights)
new_layer_vars = [var for var in tf.global_variables() if var.name.startswith('fc2/')]
init_new_vars = tf.variables_initializer(new_layer_vars)

# Saver to load only original model weights (exclude FC2 to prevent loading errors)
saver = tf.train.Saver(var_list=[var for var in tf.global_variables() if not var.name.startswith('fc2/')])

with tf.Session() as sess:
    # Initialize new layer weights first
    sess.run(init_new_vars)
    # Load pre-trained weights from your saved model path
    saver.restore(sess, './path/to/your/saved/model')  # Replace with your actual model path
    
    # Now you can train the new FC2 layer (only it will update during training)
    # Example training setup (adjust labels to match your dataset)
    labels = tf.placeholder(tf.float32, shape=[None, 2])
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels, logits=new_output))
    optimizer = tf.train.AdamOptimizer(1e-3)
    train_op = optimizer.minimize(loss)
    
    # Run your training loops here...

Key Notes #

• Layer Names Are Critical: Make sure the original layer names exactly match what was saved—if they don't, TensorFlow won't be able to map the pre-trained weights correctly.
• Verify Freezing: Double-check that all original layers are set to trainable=False; otherwise, your pre-trained weights might get overwritten during training.
• Rewire Carefully: Don't skip reconnecting the output layer to the new FC2 layer—if you do, the new layer won't affect the model's output at all.

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