使用TensorRT优化自定义600x600分辨率SSD Inception模型报错求助

Let's break down why you're hitting that assertion failure and walk through how to fix it:

Root Cause #

The error numPriors * numLocClasses * 4 == inputDims[param.inputOrder[0]].d[0] comes from the TensorRT DetectionOutput (NMS) plugin. This check verifies that the dimension of your model's location predictions matches the total number of anchor boxes (priors) configured in the plugin.

When switching from 300x300 to 600x600 input resolution:

• Each feature map in your SSD model doubles in size (e.g., 19x19 → 38x38 for the first feature layer)
• This drastically changes the total number of anchor boxes generated by your model
• But the TensorRT plugin and UFF conversion config are still using 300x300-era anchor counts, creating a critical dimension mismatch

Step-by-Step Solution #

1. Validate Your Trained Model's Anchor Configuration

First, confirm your training pipeline.config was fully adjusted for 600x600:

• Ensure the image_resizer is set to output 600x600 images
• Update the anchor_generator section to match new feature map sizes. For SSD Inception v2 600x600, feature map dimensions should be [38, 20, 10, 6, 4, 2], with corresponding strides and anchor aspect ratios. Calculate total priors:

  Total priors = (38×38×4) + (20×20×6) + (10×10×6) + (6×6×6) + (4×4×4) + (2×2×4) = 9072
  

  (Adjust these numbers if your custom model uses different anchor counts per layer)

2. Update UFF Conversion Config (config.py)

Your config.py needs to reflect new anchor and class parameters for the UFF converter:

class UffSSDConfig:
    NUM_CLASSES = YOUR_CLASS_COUNT + 1  # +1 for background class
    IMAGE_SHAPE = (600, 600)
    # Match priors per feature layer from your training config
    NUM_PRIORS = [38*38*4, 20*20*6, 10*10*6, 6*6*6, 4*4*4, 2*2*4]
    # Match anchor sizes/ratios from your pipeline.config
    SIZES = [0.1, 0.2, 0.375, 0.55, 0.725, 0.9, 1.075]
    ASPECT_RATIOS = [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
    # Keep other default configs (NMS thresholds, etc.) as needed

Re-run the UFF conversion with this updated config:

python3 convert_to_uff.py frozen_inference_graph.pb -O NMS -p config.py

3. Modify TensorRT Sample Code (sampleUffSSD.cpp)

Update the DetectionOutput plugin parameters to match your model's anchor count and class count:

// Locate where DetectionOutputParameters are defined
DetectionOutputParameters params;
params.numClasses = YOUR_CLASS_COUNT + 1;
params.backgroundLabelId = 0;
params.topK = 200;
params.keepTopK = 100;
params.confidenceThreshold = 0.3;
params.nmsThreshold = 0.5;
params.inputOrder = {0, 1, 2};
params.isNormalized = true;
params.numPriors = 9072;  // Use the total priors you calculated earlier

Double-check that your input registration and batch stream dimensions are correctly set to 600x600 (you mentioned you already did this, just confirm no typos).

4. Recompile and Test

Recompile the sample_uff_ssd binary and run it again. If dimensions are now aligned, the assertion error should disappear.

Additional Checks #

• Use trtexec to verify your UFF model's input/output dimensions:

  trtexec --uff=frozen_graph.uff --input=Input,3x600x600 --output=NMS
  

  This will show exact output dimensions, which should match numPriors * numClasses for scores and numPriors * 4 * numClasses for boxes.
• Confirm your training process completed without errors—if the model wasn't properly trained for 600x600, even correct TensorRT configs won't fix the issue.

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