TensorFlow2 Object Detection API中CenterNet微调的Checkpoint适配方案

I get it—you’ve been following the official few-shot RetinaNet fine-tuning tutorial, but hitting shape mismatch errors when trying to adapt it to CenterNet (or EfficientDet) after changing the number of classes. The key issue here is that these models have different head structures than RetinaNet, so you need to adjust the checkpoint restoration logic to skip the class-dependent layers that don’t match your new dataset’s class count.

Let’s break this down for each model type:

Fixing CenterNet Fine-Tuning #

CenterNet’s prediction head has three main branches:

1. Heatmap branch: Outputs class-specific detection scores (this is where the shape mismatch happens, since COCO uses 90 classes but you need 1)
2. Offset branch: Predicts bounding box offsets (class-agnostic, safe to restore)
3. Size branch: Predicts bounding box dimensions (also class-agnostic)

To fix the checkpoint restore, we’ll create a "fake" checkpoint structure that only includes the parts of the model we want to retain weights for, omitting the heatmap branch so it initializes from scratch. Here’s the working code:

import tensorflow as tf
from object_detection.utils import config_util
from object_detection.builders import model_builder
import os

tf.keras.backend.clear_session()
print('Building CenterNet model and restoring weights for fine-tuning...', flush=True)

num_classes = 1
pipeline_config = 'models/research/object_detection/configs/tf2/centernet_hourglass104_512x512_coco17_tpu-8.config'
checkpoint_path = 'models/research/object_detection/test_data/checkpoint/ckpt-0'

# Load pipeline config and build the detection model
configs = config_util.get_configs_from_pipeline_file(pipeline_config)
model_config = configs['model']
model_config.center_net.num_classes = num_classes
model_config.center_net.freeze_batchnorm = True  # Optional but helpful for few-shot stability
detection_model = model_builder.build(model_config=model_config, is_training=True)

# Create a fake checkpoint structure to restore only non-class-dependent weights
fake_center_net_head = tf.compat.v2.train.Checkpoint(
    _offset_prediction_head=detection_model._center_net_head._offset_prediction_head,
    _size_prediction_head=detection_model._center_net_head._size_prediction_head,
    _shared_backbone=detection_model._center_net_head._shared_backbone,
    # Omit _heatmap_prediction_head to let it initialize from scratch
)

fake_model = tf.compat.v2.train.Checkpoint(
    _feature_extractor=detection_model._feature_extractor,
    _center_net_head=fake_center_net_head
)

# Restore weights and ignore missing variables (since we skipped the heatmap head)
ckpt = tf.compat.v2.train.Checkpoint(model=fake_model)
ckpt.restore(checkpoint_path).expect_partial()

# Run a dummy image through the model to initialize all variables
image, shapes = detection_model.preprocess(tf.zeros([1, 512, 512, 3]))
prediction_dict = detection_model.predict(image, shapes)
_ = detection_model.postprocess(prediction_dict, shapes)

print('CenterNet weights restored successfully!')

Fixing EfficientDet Fine-Tuning #

EfficientDet is part of the SSD family, so its structure is closer to RetinaNet—but with a few tweaks. Its box predictor has separate class and box regression heads. We just need to skip restoring the class prediction head (since its output channels depend on the number of classes). Here’s the code:

import tensorflow as tf
from object_detection.utils import config_util
from object_detection.builders import model_builder
import os

tf.keras.backend.clear_session()
print('Building EfficientDet model and restoring weights...', flush=True)

num_classes = 1
pipeline_config = 'models/research/object_detection/configs/tf2/efficientdet_d0_coco17_tpu-32.config'
checkpoint_path = 'models/research/object_detection/test_data/checkpoint/ckpt-0'

# Load config and build model
configs = config_util.get_configs_from_pipeline_file(pipeline_config)
model_config = configs['model']
model_config.ssd.num_classes = num_classes
model_config.ssd.freeze_batchnorm = True
detection_model = model_builder.build(model_config=model_config, is_training=True)

# Fake box predictor to skip class-specific weights
fake_box_predictor = tf.compat.v2.train.Checkpoint(
    _base_tower_layers_for_heads=detection_model._box_predictor._base_tower_layers_for_heads,
    _box_prediction_head=detection_model._box_predictor._box_prediction_head,
    # Omit _class_prediction_head to reinitialize it
)

fake_model = tf.compat.v2.train.Checkpoint(
    _feature_extractor=detection_model._feature_extractor,
    _box_predictor=fake_box_predictor
)

# Restore weights
ckpt = tf.compat.v2.train.Checkpoint(model=fake_model)
ckpt.restore(checkpoint_path).expect_partial()

# Initialize variables with dummy input
image, shapes = detection_model.preprocess(tf.zeros([1, 640, 640, 3]))
prediction_dict = detection_model.predict(image, shapes)
_ = detection_model.postprocess(prediction_dict, shapes)

print('EfficientDet weights restored successfully!')

Key Notes #

• Verify Model Structure: If you’re unsure about a model’s internal components, use dir(detection_model) or print specific attributes (e.g., print(dir(detection_model._center_net_head))) to see exactly what layers are available.
• expect_partial() is Normal: TensorFlow will warn about uninitialized variables—this is expected because we’re intentionally skipping the class-dependent heads.
• TF2 Configs: Make sure you’re using config files from the configs/tf2/ directory, as older TF1 configs won’t work with the TF2 API.

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