TensorFlow自定义类别目标检测：训练/导出推理图的张量形状错误

Hey there! I totally get the frustration when things work smoothly during training but break at export—let’s sort this out together.

Why This Happens #

The core issue is a shape mismatch between the pre-trained COCO model weights (built for 90 classes) and your custom 11-class setup. When you only change num_classes in the pipeline config, the training process might initialize the final classification layer for 11 classes on the fly. But when exporting, the script tries to load the original 90-class weights for that layer—hence the lhs shape= [12] rhs shape= [91] error (the extra 1 accounts for the background class).

When you switch back to 90 classes, everything matches the pre-trained weights, so export works—but that doesn’t help your custom 11-class detector, of course.

Step-by-Step Fixes #

1. Update All num_classes References in Your Pipeline Config

Don’t just change the top-level num_classes—for Faster RCNN, you need to update it in two key places:

• Top-level under model.faster_rcnn:

  model {
    faster_rcnn {
      num_classes: 11
      ...
  

• Inside the rcnn_box_predictor section (where final classification/box prediction happens):

  rcnn_box_predictor {
        ...
        num_classes: 11
        ...
      }
  

Also double-check that label_map_path in the config points directly to your 11-class label_map.pbtxt file.

2. Export Using Your Trained Checkpoint (Not the Pre-Trained COCO Weights)

You already trained the model successfully, so you should have checkpoint files in your training directory (look for model.ckpt-XXXX.data-00000-of-00001, model.ckpt-XXXX.index, model.ckpt-XXXX.meta where XXXX is your training step count).

Run the export script with these trained weights instead of the original COCO checkpoint:

python export_inference_graph.py \
  --input_type image_tensor \
  --pipeline_config_path path/to/your/modified/pipeline.config \
  --trained_checkpoint_prefix path/to/your/training_dir/model.ckpt-XXXX \
  --output_directory path/to/your/exported_model_dir

Note: Don’t add file extensions to --trained_checkpoint_prefix—just use model.ckpt-XXXX.

3. Verify Fine-Tuning Settings (If Needed)

If issues persist, check your train_config section to ensure the final layer is trained properly:

• Set fine_tune_checkpoint_type: "detection"
• Ensure from_detection_checkpoint: true
• Optional: Explicitly freeze all layers except the final classification/box predictor layers to focus training on your custom classes (though this isn’t strictly necessary if your initial training worked).

Why Switching Back to 90 Classes Works (But Isn’t Useful) #

When you set num_classes back to 90, the config matches the shape of the pre-trained COCO weights, so the export script loads them without issues. But this means your model still detects the original COCO classes, not your 11 custom ones—so it’s just a temporary workaround, not a real solution.

Give these steps a try, and your 11-class detector should export successfully!

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