微调SSD模型时，类别数变更后TF OD API的分类权重处理逻辑

Great question! When fine-tuning an SSD model with the TensorFlow Object Detection API while aiming to reuse not just the feature extractor weights but also the classification/box regression heads, here's how the API handles classification weight tensors when your new model's class count differs from the pre-trained checkpoint:

The API relies on tensor name and shape matching to load weights from a pre-trained checkpoint. For SSD models, this plays out differently for classification vs. regression heads:

• Classification Head Weights
  SSD's classification head (per-feature-layer convolutional layers) outputs a tensor with shape [num_anchors * (num_classes + 1)] (the +1 accounts for the background class). If your new model has a different number of classes, the final dimension of this weight tensor won't match the pre-trained version.

  In this case, the API will automatically skip loading these mismatched classification weight tensors. Instead, it initializes the new classification weights from scratch using the default initializer specified in your pipeline config (typically Xavier or random normal initialization).

• Box Regression Head Weights
  The regression head outputs [num_anchors * 4] values (for the 4 bounding box offset coordinates per anchor), which is independent of class count. As long as your new model uses the same anchor configuration as the pre-trained checkpoint, these weights will load successfully without issues.

If you want to reuse parts of the pre-trained classification weights (e.g., for overlapping classes between the pre-trained dataset and your custom dataset), you'll need to implement a custom weight loading workflow, since the API doesn't support this out of the box. Here's a high-level approach:

1. Inspect the pre-trained checkpoint
   First, list all variables in the checkpoint to find the exact names of the classification weights:

   import tensorflow as tf
   pretrained_ckpt_path = "path/to/your/pretrained/checkpoint"
   for var_name, var_shape in tf.train.list_variables(pretrained_ckpt_path):
       if "ClassPredictor" in var_name:  # Adjust this to match your SSD variant's naming
           print(f"Variable: {var_name}, Shape: {var_shape}")
   

2. Load and modify the pre-trained weights
   Extract the relevant parts of the pre-trained classification weights, then merge them with randomly initialized weights for your new classes:

   num_anchors = 9  # Adjust to match your SSD's anchor count per feature layer
   old_classes_plus_bg = 91  # e.g., COCO's 90 classes + 1 background
   new_classes_plus_bg = 5  # e.g., your 4 custom classes + 1 background
   overlap_class_count = 2  # Number of classes shared between pre-trained and custom dataset
   
   # Load pre-trained classification weights
   pretrained_cls_weights = tf.train.load_variable(pretrained_ckpt_path, "FeatureExtractor/Conv/ClassPredictor/weights")
   pretrained_cls_bias = tf.train.load_variable(pretrained_ckpt_path, "FeatureExtractor/Conv/ClassPredictor/biases")
   
   # Extract weights for background and overlapping classes
   overlap_weights = pretrained_cls_weights[..., :num_anchors * overlap_class_count]
   overlap_bias = pretrained_cls_bias[:num_anchors * overlap_class_count]
   
   # Initialize weights for new classes
   new_cls_weights = tf.random.normal(
       shape=pretrained_cls_weights.shape[:-1] + (num_anchors * (new_classes_plus_bg - overlap_class_count),)
   )
   new_cls_bias = tf.random.normal(shape=(num_anchors * (new_classes_plus_bg - overlap_class_count),))
   
   # Merge weights and biases
   merged_cls_weights = tf.concat([overlap_weights, new_cls_weights], axis=-1)
   merged_cls_bias = tf.concat([overlap_bias, new_cls_bias], axis=0)
   
   # Assign merged weights to your new model
   your_model.get_layer("ClassPredictor").set_weights([merged_cls_weights.numpy(), merged_cls_bias.numpy()])
   

3. Proceed with fine-tuning
   After manually setting these modified weights, you can continue with the standard fine-tuning workflow using the Object Detection API.

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