如何在TensorFlow目标检测ms_coco_model中添加新类并训练

我之前也做过类似的COCO模型扩展任务，结合你之前用单类浣熊数据集的经验，按以下步骤操作就能顺利完成新类的添加和训练：

1. 准备包含新类的数据集 #

首先要把你的新类别数据（图片+标注）转换成TensorFlow Object Detection API支持的TFRecord格式。这里要注意：

• 新类别要分配唯一的类别ID（比如91，因为COCO原有90类的ID是1-90），绝对不能和原有类别ID重复。
• 如果想要保留原COCO类的检测能力，最好在训练数据中混入一部分COCO原数据集的样本；如果只需要新类的检测但不破坏原类性能，也可以只训练新类样本，但要配合后面的冻结层设置。
• 标注格式要和原COCO保持一致，确保TFRecord里的每个样本标注都包含正确的class_id字段。

2. 修改模型配置文件（核心步骤） #

基于ssd_mobilenet_v1_coco的原始配置文件，做以下关键修改：

• 调整类别数：找到num_classes参数，将原来的90改为91（原90类+你的新类）。
• 指定预训练Checkpoint：确认fine_tune_checkpoint字段指向你下载的ssd_mobilenet_v1_coco_checkpoint/model.ckpt，同时设置from_detection_checkpoint: true，确保正确加载预训练检测模型。
• 冻结特征提取层：为了保留原COCO模型的通用特征提取能力，避免破坏原有类的检测效果，建议冻结特征提取器层，只训练检测头和新类相关的分支。在train_config中添加：

  fine_tune_checkpoint_type: "detection"
  freeze_variables:
    - "FeatureExtractor"
  

• 更新输入路径：把train_input_reader和eval_input_reader中的input_path改成你新生成的TFRecord文件路径，label_map_path指向更新后的标签映射文件。

3. 更新标签映射文件 #

复制原COCO的label_map.pbtxt（包含1-90类的定义），在文件末尾添加你的新类别，格式如下：

item {
  id: 91
  name: '你的新类别名称'
}

务必保证ID是91（或大于90的唯一值），名称要和你的标注中的类别名完全一致。

4. 启动训练 #

用官方提供的训练脚本启动训练，这里以TF2版本为例，命令如下：

python model_main_tf2.py \
  --model_dir=./my_new_class_training \
  --pipeline_config_path=./modified_ssd_mobilenet_v1_coco.config

• 建议设置合适的训练步数：因为是微调预训练模型，不需要太多步数，比如num_train_steps: 20000、num_eval_steps: 1000就足够。
• 训练过程中可以用TensorBoard监控损失：tensorboard --logdir=./my_new_class_training，如果损失稳定下降，说明训练正常。

5. 验证与导出模型 #

训练完成后，导出冻结的推理图，用测试图片同时验证原COCO类和新类的检测效果，确保两者都能被正确识别。如果原类检测效果下降，可能是冻结层设置有问题，或者训练步数过多导致过拟合。

避坑提醒 #

• 别搞混类别ID：如果新类ID和原有类重复，会导致训练时类别混淆，严重影响检测效果。
• 确保TFRecord格式正确：如果训练时出现数据加载错误，检查TFRecord的生成脚本，确认每个样本的标注都包含正确的class_id、bbox等字段。
• 不要冻结过多层：如果冻结了检测头的层，新类的训练效果会很差，只冻结FeatureExtractor层就足够。

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