如何转换.pb文件中的TensorFlow算子？Mobilenet-v1适配300x300输入改造求助

Got it, let's work through how to swap that fixed-size 7x7 avg_pooling in your MobileNet-V1 .pb model with a global reduce_mean so it can handle 300x300 inputs. Here's a practical, step-by-step method using TensorFlow tools:

Step 1: Load the Original Model and Locate the Target AvgPool Node #

First, we need to load the .pb file and find the specific 7x7 average pooling node. This code will help you identify it:

import tensorflow as tf
from tensorflow.core.framework import graph_pb2

# Load the pre-trained .pb model
with open('mobilenet-v1.pb', 'rb') as f:
    graph_def = graph_pb2.GraphDef()
    graph_def.ParseFromString(f.read())

# Search for the 7x7 AvgPool node
target_avg_pool_node = None
for node in graph_def.node:
    # Check if the node is AvgPool with kernel size [1,7,7,1]
    if node.op == 'AvgPool' and node.attr['ksize'].list.i == [1, 7, 7, 1]:
        target_avg_pool_node = node
        print(f"Found target pooling node: {node.name}")
        break

if not target_avg_pool_node:
    raise ValueError("Could not find the 7x7 AvgPool node in your model")

Step 2: Create the ReduceMean Node #

We'll build a new ReduceMean node that performs global average pooling (works with any input size) instead of fixed-window pooling. This node will take the same input as the original AvgPool node:

# Create a new ReduceMean node for global average pooling
reduce_mean_node = graph_pb2.NodeDef()
reduce_mean_node.name = 'global_avg_pool'  # Name your new node appropriately
reduce_mean_node.op = 'ReduceMean'

# Use the same input as the original AvgPool node
reduce_mean_node.input.extend([target_avg_pool_node.input[0]])

# Configure the node to average over height and width dimensions (axes 1 and 2)
axis_attr = reduce_mean_node.attr['axis']
axis_attr.list.i.extend([1, 2])

# Keep dimensions to match the original AvgPool output shape (avoids breaking downstream layers)
keep_dims_attr = reduce_mean_node.attr['keep_dims']
keep_dims_attr.b = True

Step 3: Replace the AvgPool Node and Update Graph References #

Next, we need to remove the original AvgPool node from the graph, add our new ReduceMean node, and update all downstream nodes that referenced the old pooling node:

# Filter out the original AvgPool node
modified_nodes = [node for node in graph_def.node if node.name != target_avg_pool_node.name]
# Add the new ReduceMean node to the graph
modified_nodes.append(reduce_mean_node)

# Update all nodes that used the AvgPool node as input
for node in modified_nodes:
    updated_inputs = []
    for inp in node.input:
        if inp == target_avg_pool_node.name:
            updated_inputs.append(reduce_mean_node.name)
        else:
            updated_inputs.append(inp)
    node.input[:] = updated_inputs

# Assign the modified nodes back to the graph definition
graph_def.node[:] = modified_nodes

Step 4: Save the Modified Model #

Now save the adjusted graph to a new .pb file:

# Write the modified graph to a new file
with open('mobilenet-v1-modified.pb', 'wb') as f:
    f.write(graph_def.SerializeToString())

print("Modified model saved as 'mobilenet-v1-modified.pb'")

Step 5: Verify the Modified Model #

Test that the new model works with 300x300 inputs. Note: You may need to adjust the input/output tensor names to match your specific model:

# Test inference with a 300x300 input
tf.compat.v1.reset_default_graph()
with tf.compat.v1.Session() as sess:
    tf.import_graph_def(graph_def, name='')
    
    # Replace these with your model's actual input/output node names
    input_tensor = sess.graph.get_tensor_by_name('input:0')
    output_tensor = sess.graph.get_tensor_by_name('MobilenetV1/Predictions/Reshape_1:0')
    
    # Create a test 300x300 input tensor
    test_input = tf.random.normal([1, 300, 300, 3])
    output = sess.run(output_tensor, feed_dict={input_tensor: test_input.numpy()})
    
    print(f"Success! Output shape: {output.shape}")
    # The output shape should match the original model's output (e.g., [1, 1000] for ImageNet)

Key Notes #

• Node Names: Use TensorBoard or print all node names (loop through graph_def.node) to confirm your model's input/output and intermediate node names if the example names don't match.
• TensorFlow Version: This code works with both TensorFlow 1.x and 2.x (using compatibility mode for GraphDef handling).
• Keep Dims: Setting keep_dims=True ensures the ReduceMean output has the same shape as the original AvgPool, so downstream layers (like fully connected layers) don't need adjustments.

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