如何基于SSD Mobilenet COCO V1扩展自定义类并保留原有90类

我来帮你搞定这个需求——在SSD MobileNet COCO V1基础上保留原有90类检测能力，同时添加自定义的man和woman类别。这本质是增量迁移学习的场景，操作起来不算复杂，但几个关键细节得盯紧：

1. 数据集与类别映射准备 #

首先得把数据和类别体系统一好：

• 把你的man、woman标注数据转换成TensorFlow Object Detection API支持的TFRecord格式（如果还没转，用API里的create_pascal_tf_record.py或者自定义脚本都可以）。
• 新建一个label map文件，把原COCO的90类（ID 1-90）和你的2类（比如man设为91，woman设为92）整合进去，格式示例：

  item {
    id: 1
    name: 'person'
  }
  # ... 这里省略原COCO的90类定义
  item {
    id: 91
    name: 'man'
  }
  item {
    id: 92
    name: 'woman'
  }
  

• 如果你想最大化保留原模型的90类能力，建议把原COCO的TFRecord和你自己的自定义TFRecord一起作为训练数据（不用手动合并文件，后续在配置里指定多个路径即可）。

2. 修改模型配置文件 #

找到SSD MobileNet V1 COCO的原始配置文件（比如ssd_mobilenet_v1_coco.config），做这几个关键修改：

• 调整类别总数：把num_classes从90改成92（90+2），这是核心，不然模型输出维度会不匹配。
• 指定新的label map：在train_input_reader和eval_input_reader里，把label_map_path指向你刚创建的92类label map文件。
• 设置预训练权重：在train_config里，fine_tune_checkpoint设为SSD MobileNet V1 COCO预训练模型的ckpt路径（比如ssd_mobilenet_v1_coco_2018_01_28/model.ckpt），同时设置fine_tune_checkpoint_type: "detection"，这样会加载整个检测模型的权重，而非只加载特征提取器。
• 微调策略：为了不破坏原模型对90类的检测能力，建议用小学习率（比如learning_rate_base: 1e-5），也可以选择先冻结特征提取层（MobileNet部分），只训练后面的检测头，等损失稳定后再解冻全层微调。
• 指定训练数据：在train_input_reader的input_path里，把原COCO的TFRecord路径和你的自定义TFRecord路径都加进去，用列表形式：

  input_path: "path/to/coco_train.record"
  input_path: "path/to/custom_train.record"
  

3. 启动增量训练 #

用Object Detection API的训练脚本启动训练（TF2版本用model_main_tf2.py，TF1用train.py）：

# TF2版本示例
python model_main_tf2.py \
  --model_dir=./my_custom_ssd_model \
  --pipeline_config_path=./ssd_mobilenet_v1_coco_custom.config \
  --num_train_steps=50000 \
  --num_eval_steps=5000

训练过程中可以用TensorBoard监控损失曲线，确保分类损失和定位损失都在稳步下降。

4. 导出模型并测试 #

训练完成后，导出可用于推理的冻结模型：

# TF2版本示例
python exporter_main_v2.py \
  --input_type image_tensor \
  --pipeline_config_path=./ssd_mobilenet_v1_coco_custom.config \
  --trained_checkpoint_dir=./my_custom_ssd_model \
  --output_directory=./exported_custom_ssd

之后你就可以用这个导出的模型做视频实时检测了：用OpenCV读取视频帧，把帧转换成模型要求的格式（比如RGB、尺寸匹配），传入模型推理，然后把检测结果（框、类别、置信度）绘制到帧上，再输出成视频或者实时显示。

• 如果不想用原COCO数据训练，只靠自定义数据的话，原90类的检测能力大概率会下降，因为模型会“遗忘”旧类别。所以最好混合原COCO数据和自定义数据一起训练。
• 自定义数据的标注质量要高，尽量覆盖不同场景、光照、角度，和原COCO数据的分布越接近，模型泛化效果越好。
• 训练时可以先固定特征提取层，只训练检测头，等损失稳定后再解冻特征层用更小的学习率微调，这样能平衡训练效率和模型性能。

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