手把手教你：如何将 Deformable DETR 预训练权重适配到自己的3类小数据集

在目标检测领域，迁移学习已成为小数据集场景下的黄金法则。Deformable DETR 作为 DETR 系列中的佼佼者，凭借其可变形注意力机制，在保持端到端检测优势的同时，显著提升了小目标检测性能。但当官方预训练模型遇到自定义数据集时，类别数不匹配就像一副尺寸不合的手套——直接使用会导致关键层维度冲突。本文将深入解析权重文件结构，演示如何通过精准手术式修改，让预训练模型完美适配你的3类数据集。

1. 环境配置与核心原理

1.1 环境准备

工欲善其事，必先利其器。以下是经过验证的环境组合：

# 基础环境
conda create -n deformable_detr python=3.7
conda install pytorch==1.11.0 torchvision==0.12.0 cudatoolkit=11.3 -c pytorch

注意：PyTorch 1.10.0 存在符号链接错误，务必使用1.11.0版本。若遇到 MultiScaleDeformableAttention 编译错误，可尝试以下修复方案：

# 重新编译注意力模块
cd ./models/ops
rm -rf build
sh ./make.sh

1.2 模型架构关键点

Deformable DETR 的类别适配核心在于两个关键层：

层名称	原始维度	3类数据集所需维度	作用描述
class_embed	[91, 256]	[4, 256]	分类头预测（含背景类）
query_embed	[300, 512]	[50, 512]	可学习的位置查询向量

表：关键层维度对照表（COCO→3类数据集）

# 原始维度查看代码示例
pretrained_weights = torch.load('r50_deformable_detr-checkpoint.pth')
print(pretrained_weights['model']['class_embed.0.weight'].shape)  # 输出torch.Size([91, 256])

2. 权重文件手术指南

2.1 权重矩阵改造实战

以下完整脚本实现自动化适配：

import torch

def adapt_weights(input_path, output_path, num_classes=3, num_queries=50):
    weights = torch.load(input_path)
    
    # 分类头改造（6个连续层）
    for i in range(6):
        weights['model'][f'class_embed.{i}.weight'] = weights['model'][f'class_embed.{i}.weight'][:num_classes+1]
        weights['model'][f'class_embed.{i}.bias'] = weights['model'][f'class_embed.{i}.bias'][:num_classes+1]
    
    # 查询向量改造
    if 'query_embed.weight' in weights['model']:
        weights['model']['query_embed.weight'] = weights['model']['query_embed.weight'][:num_queries]
    
    torch.save(weights, output_path)
    print(f"适配后的权重已保存至 {output_path}")

# 使用示例
adapt_weights(
    input_path='r50_deformable_detr-checkpoint.pth',
    output_path='deformable_detr-r50_3.pth'
)

2.2 维度修改原理详解

1. 分类头改造逻辑：

   • 原始COCO权重包含80个物体类+1个背景类
   • 3类数据集需要保留前3个物体类+1个背景类
   • resize_操作直接修改张量存储结构，比切片更高效

2. 查询向量优化技巧：

   • 原始300个查询可能造成计算浪费
   • 根据实际需求调整num_queries参数（建议值50-100）

3. 配置文件联动调整

3.1 关键参数对照表

文件路径	需要修改的参数	示例值
configs/r50_deformable_detr.sh	--num_queries	50
models/deformable_detr.py	num_classes	3
main.py	--dataset_file	'custom'

表：必须同步修改的配置文件参数

3.2 训练启动命令优化

# 单卡训练示例（适配修改后的配置）
GPUS_PER_NODE=1 \
./tools/run_with_submitit.py \
    --config-file configs/r50_deformable_detr.sh \
    --num-gpus 1 \
    --num-queries 50 \
    --output-dir exps/r50_deformable_detr_3class

提示：首次运行建议添加--eval参数验证配置正确性

4. 实战调试与效果验证

4.1 常见错误排查指南

• 维度不匹配错误：

  # 典型报错示例
  RuntimeError: size mismatch, m1: [32 x 256], m2: [91 x 256]
  

  解决方案：检查所有6个class_embed层是否全部修改

• CUDA内存不足： 降低--batch-size（默认值为2），或减少num_queries

4.2 可视化检测代码优化

# 改进的检测结果可视化（支持中文标签）
def plot_result_cn(pil_img, scores, boxes, class_names=['猫', '狗', '鸟']):
    cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    for score, (x1, y1, x2, y2) in zip(scores, boxes):
        cls_id = score.argmax()
        cv2.rectangle(cv_img, (int(x1), int(y1)), (int(x2), int(y2)), (0,255,255), 2)
        text = f"{class_names[cls_id]}:{score[cls_id]:.1%}"
        cv2.putText(cv_img, text, (int(x1), int(y1)-10), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)
    return cv_img

4.3 性能优化技巧

1. 冻结骨干网络：

   # 在build_model后添加
   for name, param in model.backbone.named_parameters():
       param.requires_grad = False
   

2. 学习率分层设置：

   --lr 1e-4 \
   --lr-backbone 0 \
   --lr-linear-proj 1e-5
   

在实际项目中，我发现当类别数从80降到3时，适当降低分类头学习率（如1e-5）能有效避免过拟合。同时将训练周期从默认的50缩减到20，即可获得不错的效果。