Jetson TX2环境下如何将SSD_MOBILENET V1冻结图转为TensorRT模型？

Hey there! Let me walk you through converting your SSD_MobileNet V1 frozen graph to a TensorRT model with your exact setup (Jetson TX2, TensorFlow 1.6, JetPack 3.2/TensorRT 3.0, and the UFF tool). I’ve done this a few times on TX2, so here’s the step-by-step breakdown that should work for you:

First, make sure you have a frozen .pb file for SSD_MobileNet V1—this is the model where all variables are converted to constants. If you haven’t generated it yet, use TensorFlow’s freeze_graph.py tool:

python freeze_graph.py \
  --input_graph=path/to/your/graph.pbtxt \
  --input_checkpoint=path/to/your/model.ckpt \
  --input_binary=false \
  --output_graph=path/to/frozen_inference_graph.pb \
  --output_node_names="detection_boxes,detection_scores,detection_classes,num_detections"

Pro tip: The output node names here are the standard ones for SSD_MobileNet V1—double-check they match your model’s actual output nodes if you’ve made custom modifications.

TensorRT 3.0 relies on UFF (Universal Framework Format) as an input, so we’ll use the UFF tool to convert the .pb file. You can do this via command line or Python API:

Option 1: Command Line (Quick & Simple) #

convert-to-uff path/to/frozen_inference_graph.pb \
  --output-file=ssd_mobilenet_v1.uff \
  --input-node-names="image_tensor" \
  --output-node-names="detection_boxes,detection_scores,detection_classes,num_detections"

image_tensor is the standard input node name for SSD_MobileNet V1—adjust if your model uses a different name.

Option 2: Python API (More Flexible) #

If you need to tweak nodes or handle edge cases, use the UFF Python API:

import uff

frozen_graph_path = "path/to/frozen_inference_graph.pb"
uff_model_path = "ssd_mobilenet_v1.uff"

# Define input/output nodes matching your model
input_nodes = ["image_tensor"]
output_nodes = ["detection_boxes", "detection_scores", "detection_classes", "num_detections"]

# Convert and save the UFF model
uff.from_tensorflow_frozen_model(
    frozen_graph_path,
    output_nodes,
    input_nodes=input_nodes,
    output_filename=uff_model_path
)

With the UFF model ready, we’ll create a TensorRT engine optimized for the TX2’s GPU. Here’s a Python example (great for quick testing):

import tensorrt as trt
import numpy as np

# Set up a TensorRT logger (suppress verbose logs with WARNING level)
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

# Model parameters (match SSD_MobileNet V1 specs)
INPUT_SHAPE = (3, 300, 300)  # CHW format: channels, height, width
INPUT_NAME = "image_tensor"
OUTPUT_NAMES = ["detection_boxes", "detection_scores", "detection_classes", "num_detections"]

def build_trt_engine(uff_model_path):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
        # Configure engine for TX2
        builder.max_workspace_size = 1 << 28  # 256MB (adjust based on TX2's available memory)
        builder.fp16_mode = True  # Enable FP16 for faster inference (TX2 supports this natively)
        builder.max_batch_size = 1

        # Register input/output nodes with the parser
        parser.register_input(INPUT_NAME, INPUT_SHAPE)
        for output_name in OUTPUT_NAMES:
            parser.register_output(output_name)
        
        # Parse the UFF model into the TensorRT network
        parser.parse(uff_model_path, network)

        # Build and return the optimized engine
        return builder.build_cuda_engine(network)

# Build the engine
engine = build_trt_engine("ssd_mobilenet_v1.uff")

# Example inference workflow (simplified)
with engine.create_execution_context() as context:
    # Generate dummy input (replace with your actual image data)
    input_data = np.random.randn(1, *INPUT_SHAPE).astype(np.float32)  # Batch size 1
    
    # Prepare output buffers
    output_buffers = [np.empty(context.get_binding_shape(i), dtype=np.float32) for i in range(engine.num_bindings)]
    
    # Bind memory (in production, use CUDA memory allocation for better performance)
    bindings = [int(input_data.ctypes.data)] + [int(buf.ctypes.data) for buf in output_buffers]
    
    # Run inference
    context.execute_v2(bindings)
    
    # Extract results
    boxes, scores, classes, num_detections = output_buffers
    print(f"Detected {int(num_detections[0])} objects")

• Node Name Accuracy: Always verify input/output node names match your frozen graph—use tf.get_default_graph().get_operations() in TensorFlow to list all nodes if you’re unsure.
• Memory Limits: TX2 has limited RAM, so don’t set max_workspace_size higher than 512MB unless you’re sure you have the space.
• FP16 Speedup: Enabling fp16_mode is a no-brainer for TX2—it cuts inference time significantly with minimal precision loss for SSD models.
• Version Compatibility: Your setup (TF 1.6 + TRT 3.0 + JetPack 3.2) is fully compatible, so you won’t hit version mismatch issues.

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