如何在GPU上运行含伪量化节点的MobileNet图并提取指定层输出？

Got it, let's walk through how to pull that last pointwise convolution output while ditching the unnecessary fc and softmax layers. Here's a step-by-step solution with working code:

Step 1: Complete the Graph Loading Function #

First, let's finish that load_graph function properly to load your frozen protobuf graph correctly:

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

def load_graph(frozen_graph_filename):
    # Load the frozen graph from disk
    with tf.io.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = graph_pb2.GraphDef()
        graph_def.ParseFromString(f.read())

    # Import the graph definition into a new TensorFlow graph
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name="")
    return graph, graph_def

Step 2: Find Your Target Layer's Tensor Name #

You need to identify the exact name of the output tensor from the last pointwise convolution. Since this is a training-phase quantized model with pseudo-quantization nodes, run this code to inspect the graph:

graph, graph_def = load_graph("your_frozen_model.pb")

# Print operations containing "pointwise" to locate the last one
for op in graph.get_operations():
    if "pointwise" in op.name.lower():
        print(f"Operation Name: {op.name}")
        print(f"Output Tensors: {[t.name for t in op.outputs]}\n")

For MobileNet, the last pointwise convolution is typically named something like Conv2d_13_pointwise. Its output might be:

• Conv2d_13_pointwise/Relu6:0 (raw activation output before pseudo-quantization)
• Conv2d_13_pointwise/Relu6/FakeQuantWithMinMaxVars:0 (quantized simulation output, matching test-time behavior)

Pick the tensor that aligns with your needs.

Step 3: Run Inference to Fetch the Target Tensor #

Once you have the tensor name, you can run inference and pull that specific output instead of the final softmax/fc layer:

def get_pointwise_output(graph, input_data, target_tensor_name):
    # Get the input tensor (usually named 'input:0' for MobileNet)
    input_tensor = graph.get_tensor_by_name("input:0")
    # Get your target output tensor
    target_tensor = graph.get_tensor_by_name(target_tensor_name)

    # Run inference
    with tf.Session(graph=graph) as sess:
        output = sess.run(target_tensor, feed_dict={input_tensor: input_data})
    return output

# Example usage:
# Replace with your preprocessed input (shape: (1, 224, 224, 3) for standard MobileNet)
input_data = preprocess_your_input(your_image)
# Use the tensor name you found earlier
target_tensor = "Conv2d_13_pointwise/Relu6/FakeQuantWithMinMaxVars:0"
pointwise_output = get_pointwise_output(graph, input_data, target_tensor)

Step 4: Optional - Prune the Graph for Deployment #

If you want a trimmed-down frozen graph that only computes up to your target layer (removing fc/softmax entirely), use this code to prune the graph:

from tensorflow.python.framework.graph_util import extract_sub_graph

def prune_to_target_layer(graph_def, target_tensor_name):
    # Extract the subgraph up to your target operation
    target_op_name = target_tensor_name.split(":")[0]
    pruned_graph_def = extract_sub_graph(graph_def, [target_op_name])

    # Save the pruned graph
    with tf.io.gfile.GFile("pruned_mobilenet.pb", "wb") as f:
        f.write(pruned_graph_def.SerializeToString())
    return pruned_graph_def

# Usage:
pruned_graph_def = prune_to_target_layer(graph_def, target_tensor)

Quick Notes: #

• Ensure your input data is preprocessed correctly: Resize to 224x224, normalize with mean values [123.68, 116.779, 103.939], and adjust channel order if needed.
• For TensorFlow 2.x, replace tf.Session with tf.compat.v1.Session to maintain compatibility with frozen graphs.

内容的提问来源于stack exchange，提问作者A. Sam