深度学习 | Detectron2实战：从零构建自定义检测模型

1. Detectron2安装与环境配置Detectron2作为Facebook AI Research推出的目标检测框架其安装过程在不同操作系统下存在显著差异。对于Linux用户而言安装过程相对简单直接。首先需要确保系统满足以下基础依赖Python ≥ 3.6PyTorch ≥ 1.8及匹配版本的torchvisionOpenCV可选但可视化演示需要GCC/G ≥ 5.4CUDA和cuDNNGPU加速必需在Ubuntu系统下可通过以下命令快速安装python -m pip install githttps://github.com/facebookresearch/detectron2.git若遇到权限问题可添加--user参数。对于需要本地编译的场景建议先克隆仓库再安装git clone https://github.com/facebookresearch/detectron2.git python -m pip install -e detectron2Windows平台的安装则更为复杂需要预先安装Visual Studio 2019的C编译环境。关键步骤包括通过Visual Studio Installer添加使用C的桌面开发工作负载安装pycocotools时需特别注意依赖项pip install githttps://github.com/philferriere/cocoapi.git#subdirectoryPythonAPI推荐使用conda管理环境以避免依赖冲突conda install cython ninja pywin322. 自定义数据集准备2.1 COCO格式详解Detectron2默认支持COCO格式的数据标注其目录结构应包含dataset_root/ ├── train_images/ │ ├── 001.jpg │ └── 002.jpg ├── val_images/ │ └── 003.jpg └── annotations/ ├── instances_train.json └── instances_val.json标注JSON文件的核心结构包括{ images: [{id: 1, file_name: 001.jpg, width: 640, height: 480}], annotations: [{ id: 1, image_id: 1, category_id: 1, bbox: [x,y,width,height], area: 1200, iscrowd: 0 }], categories: [{id: 1, name: cat}] }2.2 数据集注册实战对于非COCO格式数据需自定义注册逻辑。以下示例展示遥感图像数据集的注册方法from detectron2.data import DatasetCatalog, MetadataCatalog def register_remote_sensing_dataset(): classes [{id:1, name:airplane}, {id:2, name:ship}] for split in [train, val]: DatasetCatalog.register( fremote_sensing_{split}, lambda: load_custom_json(fannotations/{split}.json) ) MetadataCatalog.get(fremote_sensing_{split}).thing_classes [ c[name] for c in classes ]2.3 数据增强策略Detectron2提供灵活的数据增强配置可通过DatasetMapper实现from detectron2.data import transforms as T augmentations [ T.RandomFlip(prob0.5), T.RandomBrightness(0.8, 1.2), T.RandomContrast(0.8, 1.2), T.Resize((800, 800)) ] mapper DatasetMapper(cfg, is_trainTrue, augmentationsaugmentations) train_loader build_detection_train_loader(cfg, mappermapper)3. 模型架构深度定制3.1 Backbone网络改造Detectron2支持自定义特征提取网络以下实现一个简化版ResNetfrom detectron2.modeling import BACKBONE_REGISTRY BACKBONE_REGISTRY.register() class SimpleResNet(nn.Module): def __init__(self, cfg, input_shape): super().__init__() self.conv1 nn.Conv2d(3, 64, kernel_size7, stride2, padding3) self.resblocks nn.Sequential( *[ResidualBlock(64) for _ in range(4)] ) def forward(self, x): return {features: self.resblocks(self.conv1(x))}3.2 检测头创新设计针对小目标检测场景可改进RPN网络from detectron2.modeling import RPN_HEAD_REGISTRY RPN_HEAD_REGISTRY.register() class DenseRPNHead(nn.Module): def __init__(self, cfg, input_shape): super().__init__() self.anchor_generator build_anchor_generator(cfg) self.conv nn.Conv2d(input_shape.channels, 256, 3) self.object_logits nn.Conv2d(256, len(self.anchor_generator.sizes), 1) def forward(self, features): pred_logits [] for x in features: t F.relu(self.conv(x)) pred_logits.append(self.object_logits(t)) return pred_logits4. 训练流程优化技巧4.1 学习率调度策略工业质检场景常采用渐进式学习率调整from detectron2.solver import WarmupCosineLR def build_lr_scheduler(cfg, optimizer): return WarmupCosineLR( optimizer, cfg.SOLVER.MAX_ITER, warmup_factor0.001, warmup_iters1000, warmup_methodlinear )4.2 自定义评估指标针对特定业务需求添加mAP0.5:0.95指标from detectron2.evaluation import COCOEvaluator class CustomEvaluator(COCOEvaluator): def _eval_predictions(self, predictions): super()._eval_predictions(predictions) # 添加自定义指标计算逻辑 self._results[AP_0.5_0.95] calculate_area_under_curve()4.3 混合精度训练通过Apex库启用混合精度训练加速from apex import amp model build_model(cfg) optimizer build_optimizer(cfg, model) model, optimizer amp.initialize(model, optimizer, opt_levelO1) trainer DefaultTrainer(cfg) trainer.model model trainer.optimizer optimizer5. 工业级部署方案5.1 TorchScript导出将训练好的模型转换为生产环境可用的格式from detectron2.export import scripting model build_model(cfg) scripted_model scripting.export_scripting(model, (320, 320)) torch.jit.save(scripted_model, model.pt)5.2 TensorRT加速使用TensorRT进行推理优化from detectron2.export import Caffe2Tracer tracer Caffe2Tracer(cfg, model, (320, 320)) caffe2_model tracer.export_caffe2() onnx_model export_onnx(caffe2_model)在实际工业质检项目中这种端到端的解决方案可将推理速度提升3-5倍。我曾在一个PCB缺陷检测项目中通过模型量化将部署在Jetson Xavier上的推理耗时从120ms降至28ms同时保持98.7%的检测准确率。