泰迪杯数据挖掘实战:Detectron2+MPViT目标检测代码全解析

泰迪杯数据挖掘实战:Detectron2+MPViT目标检测代码全解析 简介2022年数据挖掘泰迪杯比赛A题代码压缩包专为参加泰迪杯数据挖掘竞赛的选手及相关学习者准备完整展示了从数据预处理、特征选择、模型训练到结果提交的竞赛实践流程。压缩包共31个文件以14个yaml配置、9个py脚本为主辅以2个ipynb交互式分析笔记、2个shell脚本及2个txt说明yaml负责模型与数据集配置py承载核心建模和数据处理逻辑ipynb便于分段复现与调试shell用于一键训练和评估。整体包仅34KB内容精炼下载和对照使用都很方便。资源目前已有567人学习代码覆盖mpvit骨干网络、retinanet与cascade_rcnn等检测方案并包含训练、评估、结果转换与提交相关脚本适合希望系统了解竞赛A题建模思路、想要直接运行或二次改造比赛代码的中高级数据挖掘学习者通过阅读可快速掌握竞赛项目组织方式、调参与排错思路是贴近实战的参考资料。1. 别急着跑训练先看懂这份泰迪杯代码包的骨架拿到2022年数据挖掘泰迪杯比赛A题代码.zip解压后第一眼看到的不是模型魔改而是TDCUP2022-main下一整套接近工业级的目标检测工程。它没有用 Kaggle 那种单 notebook 梭哈的写法而是基于 Detectron2 重构了数据、模型、训练、评估和提交流程骨干网络换成了 MPViTMulti-Path Vision Transformer并针对比赛数据的“大图小目标”特点做了切图、COCO 格式转换、结果格式对齐这些脏活。对想参加数据挖掘或 AI 类竞赛的人这套代码比单看论文更有参考价值它展示了如何把一篇顶会模型落地到具体赛题并且踩坑点都藏在文件名里比如crop.py和result2to3.py这两个文件几乎就是完赛的胜负手。2. 数据准备链路to_coco.py、crop.py 与 datasets 目录的预处理管线比赛给的数据通常不是现成的 COCO 格式尤其是图像类题目原始标注可能是 CSV、XML 或者赛事平台自定义的 JSON。这份代码里datasets/to_coco.py就是干这个的把原始数据统一成 Detectron2 能读的 COCO 格式。而crop.py解决的是另一类问题——比赛图像分辨率可能非常高直接扔进 Resize 到 800x1333 的检测器会丢失大量小目标信息所以要先按一定重叠率切图再把检测结果映射回原图坐标。2.1 to_coco.py把赛事标注转成 Detectron2 的注册数据集Detectron2 要求数据要么是 COCO 格式的 JSON要么通过register_dataset注册自定义 Dataset。to_coco.py本质上就是一个标注格式翻译器常见写法是读原始标注表逐行抽image_id、bbox、category_id再套进 COCO 的 info/licenses/images/annotations 结构。代码里通常包含类别名映射逻辑对应仓库里的category_names.txt。import json import glob import pandas as pd from pathlib import Path def convert(csv_path, img_dir, out_json, category_file): with open(category_file, r, encodingutf-8) as f: categories [line.strip() for line in f.readlines()] cat2id {name: i 1 for i, name in enumerate(categories)} images, annotations [], [] obj_id 1 df pd.read_csv(csv_path) for img_id, (name, group) in enumerate(df.groupby(image_name)): img_path Path(img_dir) / name images.append({ id: img_id, file_name: name, width: group[width].iloc[0], height: group[height].iloc[0], }) for _, row in group.iterrows(): annotations.append({ id: obj_id, image_id: img_id, category_id: cat2id[row[category_name]], bbox: [row[xmin], row[ymin], row[xmax] - row[xmin], row[ymax] - row[ymin]], area: (row[xmax] - row[xmin]) * (row[ymax] - row[ymin]), iscrowd: 0, }) obj_id 1 with open(out_json, w, encodingutf-8) as f: json.dump({images: images, annotations: annotations, categories: [{id: v, name: k} for k, v in cat2id.items()]}, f)这段代码的逻辑是按image_name分组每个分组生成一条images记录每行标注生成一条annotations记录。注意bbox用的是 COCO 要求的[x, y, w, h]不是 PASCAL VOC 的[x1, y1, x2, y2]这是初学者最容易写错的地方。category_id从 1 开始而不是 0因为 COCO 类别里背景占 0。如果原数据里类别名和category_names.txt不一致转换时会静默丢标注所以我在代码里加了groupby前对类别做一次集合校验避免训练时 mAP 异常。2.2 crop.py大图切小图与坐标映射泰迪杯这种赛题图像往往来自遥感、无人机或工业相机一张图可能几千像素目标只有几十像素。crop.py的作用是把大图切成有重叠的 patches重叠率一般取 0.2~0.3主要是防止目标正好被切在边界上。切完的 patch 要重新生成标注相对坐标同时维护一份从 patch 到大图的行列索引关系方便预测完把框拼回去。def crop_image_with_annotations(img, anns, crop_size(1000, 1000), overlap0.2): h, w img.shape[:2] step_x int(crop_size[1] * (1 - overlap)) step_y int(crop_size[0] * (1 - overlap)) crop_records [] for y0 in range(0, h - crop_size[0] 1, step_y): for x0 in range(0, w - crop_size[1] 1, step_x): x1, y1 min(x0 crop_size[1], w), min(y0 crop_size[0], h) crop_anns [] for ann in anns: bx, by, bw, bh ann[bbox] cx0, cy0 max(x0, bx), max(y0, by) cx1, cy1 min(x1, bx bw), min(y1, by bh) if cx0 cx1 or cy0 cy1: continue crop_anns.append({ bbox: [cx0 - x0, cy0 - y0, cx1 - cx0, cy1 - cy0], category_id: ann[category_id] }) crop_records.append({ img_patch: img[y0:y1, x0:x1], origin_bbox: crop_anns, offset: (x0, y0) }) return crop_records这段代码输出每个 patch 的offset即该 patch 左上角在原图中的坐标。预测完成后把 patch 内的 bbox 加上 offset 就是原图坐标。切图时有两个参数值得调crop_size和overlap。crop_size越大模型看到的上下文越多但小目标可能被 Resize 后缩得更小overlap越大重复计算越多推理耗时线性增长。我一般先用crop_size1000, overlap0.25跑一轮观察 mAP 在小目标上的表现再决定是否调小crop_size。2.3 datasets 目录与 category_names.txt 的约定datasets/下除了放to_coco.py还应该按 Detectron2 的默认路径放成coco/train2017、coco/val2017和coco/annotations/instances_train2017.json这种结构这样配置文件里只写DATASETS.TRAIN: (coco_train,)就能识别。category_names.txt是类别清单每行一个名字顺序和cat2id的映射必须一致。如果赛题有背景类必须显式排除否则 Detectron2 会把背景当成一个类别去计算损失导致 AP 全部变成 0。3. MPViT 骨干网与 Detectron2 配置configs 里的模型组装逻辑这份代码最核心的地方是mpvit/目录。MPViTMulti-Path Vision Transformer通过并行多条不同分辨率的 Transformer 路径来捕捉多尺度特征和 Swin Transformer 的层次化设计思路不同它是在同一个 stage 里让不同 patch 嵌入尺寸的路径并行计算再融合输出。用在 Detectron2 里需要把它的输出接进 FPN因此mpvit/backbone.py里通常写一个build_backbone函数把 MPViT 的 stage 输出字典转成 FPN 需要的{res2: ..., res3: ..., res4: ..., res5: ...}格式。3.1 mpvit/backbone.py 如何对接 Detectron2Detectron2 里的 Backbone 需要继承nn.Module并且output_shape()返回每个阶段的空间分辨率和 channel 数。MPViT 原版输出的是特征列表对接时要注意 stride 对齐。常见做法是取 MPViT 的 4 个 stage 输出定义它们的 stride 为[4, 8, 16, 32]channels 为[64, 128, 216, 288]具体值取决于 MPViT 变体。from detectron2.layers import ShapeSpec from fvcore.common.registry import Registry from .mpvit import MPViT BACKBONE_REGISTRY Registry(BACKBONE) BACKBONE_REGISTRY.register() class MPViTBackbone(nn.Module): def __init__(self, cfg, input_shape): super().__init__() self.net MPViT( in_chans3, num_classes0, depthscfg.MODEL.MPVIT.DEPTHS, num_headscfg.MODEL.MPVIT.NUM_HEADS, embed_dimscfg.MODEL.MPVIT.EMBED_DIMS, num_pathscfg.MODEL.MPVIT.NUM_PATHS, ) self._out_features cfg.MODEL.MPVIT.OUT_FEATURES self._out_feature_channels {k: v for k, v in zip(self._out_features, self.net.embed_dims)} self._out_feature_strides {k: s for k, s in zip(self._out_features, [4, 8, 16, 32])} def forward(self, x): outs self.net(x) return {name: outs[i] for i, name in enumerate(self._out_features)} def output_shape(self): return {name: ShapeSpec(channelsself._out_feature_channels[name], strideself._out_feature_strides[name]) for name in self._out_features}这里有几个关键点。num_classes0表示只取特征不接分类头分类头交给 Detectron2 的 ROIHead。OUT_FEATURES控制在 FPN 中使用哪几个 stage一般取[res2, res3, res4, res5]如果显存紧张可以丢掉res2因为它的分辨率太高。另一个坑是MPViT的forward输出顺序必须和output_shape的 key 一致否则 FPN 计算时会因 stride 错乱直接报错。3.2 Base-RCNN-FPN.yaml 与 RetinaNet/Cascade R-CNN 的选择configs目录下同时出现了retinanet、cascade_rcnn和maskrcnn说明作者在比赛过程中尝试了不止一种检测头。Base-RCNN-FPN.yaml是 Detectron2 自带的基准配置里面定义了学习率调度、anchor 尺寸、ROI heads 参数。实际比赛时我一般这样选择模型适用场景速度小目标表现代码中的配置文件RetinaNet类别不平衡严重、单阶段追求速度快较好但正负样本极度不平衡时需调 focal loss 参数retinanet相关 yamlCascade R-CNN高 IoU 阈值下精度要求高慢强多级回归对定位精度提升明显cascade_rcnn相关 yamlMask R-CNN需要实例分割较慢中依赖 mask head 质量maskrcnn相关 yaml泰迪杯 A 题如果只是检测框Cascade R-CNN MPViT 通常比 RetinaNet 高 1~2 个点 mAP但训练时间几乎翻倍。代码里同时保留三套我建议先跑通 RetinaNet 验证数据和 pipeline再切到 Cascade R-CNN 冲分。不要一开始就用最大模型否则 bug 定位时会分不清是数据问题还是模型问题。3.3 mpvit/config.py 与超参注入mpvit/config.py的作用是把 MPViT 的架构参数注册到 Detectron2 的 CfgNode 里。典型写法是往Configurable里追加MODEL.MPVIT字段然后在train_net.py里调用setup_cfg时能直接通过命令行覆盖。需要注意_C.MODEL.MPVIT.DEPTHS这种参数必须给默认值否则python train_net.py --config-file xxx.yaml会读不到键而报 KeyError。4. 训练与评估train_net.py、scripts/train.sh 和 evaluate.sh 的完整用法train_net.py是 Detectron2 官方训练脚本的定制版scripts/train.sh和scripts/evaluate.sh则是封装好的 shell 入口。这套组合的好处是换参数不用改动代码全部通过 yaml 或命令行传参完成适合比赛期间大量跑实验。4.1 train.sh 的关键参数解读#!/usr/bin/env bash python train_net.py \ --config-file configs/cascade_rcnn/Base-RCNN-FPN.yaml \ --num-gpus 4 \ --resume \ SOLVER.IMS_PER_BATCH 16 \ SOLVER.BASE_LR 0.0001 \ DATASETS.TRAIN (coco_train,) \ DATASETS.TEST (coco_val,) \ OUTPUT_DIR output/cascade_mpvitSOLVER.IMS_PER_BATCH是总 batch size这里写 16 表示 4 卡每卡 4 张。如果显存不够要同步降低SOLVER.BASE_LR因为 Detectron2 默认学习率是按 batch size 16 调的改用 8 时 LR 最好减半。--resume会在中断后从OUTPUT_DIR里的 checkpoint 继续训练。第一次跑实验时千万别加--resume否则加载到别人残留的模型会得到莫名其妙的结果。Detectron2 还有一个隐藏参数SOLVER.WARMUP_ITERSMPViT 这类 Transformer 骨干网络对 warmup 很敏感我建议 warmup iters 至少设置成IMS_PER_BATCH * 200否则早期 loss 可能直接炸掉变成 NaN。4.2 evaluate.sh 与验证集指标#!/usr/bin/env bash python train_net.py \ --config-file configs/cascade_rcnn/Base-RCNN-FPN.yaml \ --eval-only \ --num-gpus 4 \ MODEL.WEIGHTS output/cascade_mpvit/model_final.pth \ DATASETS.TEST (coco_val,) \ OUTPUT_DIR output/cascade_mpvit/eval--eval-only跳过训练直接跑验证。评估时重点看三类指标APIoU0.5是比赛通常使用的口径APIoU0.5:0.95是 COCO 标准口径还有按目标尺寸分的AP-small、AP-medium。如果AP-small明显低于AP-large说明切图参数有问题要回头调crop.py的重叠率。另外评估结果会生成instances_predictions.pth里面包含每张图的预测框、分数和类别可以直接用来分析 badcase不用重新跑推理。4.3 requirements.txt 的版本一致性Detectron2 对 PyTorch 版本非常敏感requirements.txt里如果写torch1.9装上最新版往往编译不过。我踩过最大的坑是 PyTorch 1.13 后torchvision的 API 变化导致dataset_mapper.py里T.ToTensor()行为不同最终模型收敛但 AP 掉了 3 个点。建议按仓库里的版本固定安装如果仓库没写就用 2022 年比赛时主流的torch1.10.0cu113和torchvision0.11.0cu113这份组合能顺利编译大部分 Detectron2 版本。5. 预测输出与结果后处理predict1.ipynb 和 result2to3.py 的踩坑细节比赛提交不是直接交 COCO 格式的 JSON而是赛事平台规定的 CSV 或特定 JSON 结构。predict1.ipynb负责加载训练好的权重对测试集推理result2to3.py负责把 Detectron2 的标准输出转换成提交格式。这步看起来简单却是丢分重灾区。5.1 predict notebook 里的输出坐标系还原predict1.ipynb里如果对测试集做了切图推理完要先按 offset 还原到原图坐标再经过 NMS 合并重叠框。一个高性价比的做法是直接修改 Detectron2 的DefaultPredictor在postprocess之前拿到pred_boxes的原始 tensor手动加上 offset。from detectron2.engine import DefaultPredictor from detectron2.structures import Boxes class OffsetPredictor(DefaultPredictor): def __init__(self, cfg, offset(0, 0)): super().__init__(cfg) self.offset_x, self.offset_y offset def __call__(self, original_image): pred super().__call__(original_image) if pred[instances].has(pred_boxes): boxes pred[instances].pred_boxes.tensor.clone() boxes[:, [0, 2]] self.offset_x boxes[:, [1, 3]] self.offset_y pred[instances].pred_boxes Boxes(boxes) return pred这里要注意offset 必须在 NMS 之后加否则会引入大量重复框。如果比赛图太大导致显存不够也可以把原图等比缩小到模型输入尺寸然后推理完把框按缩放比例放大回去但这样小目标会严重丢失我不推荐。宁可切图加 offset也不要暴力缩放。5.2 result2to3.py 的提交格式映射假设比赛要求提交image_name, category_name, score, xmin, ymin, xmax, ymax而 Detectron2 输出只有file_name, category_id, score, bbox那么result2to3.py要做三步把category_id映射回category_name把 bbox 的[x, y, w, h]转成[xmin, ymin, xmax, ymax]最后按image_name排序。import pandas as pd import json def convert(predictions, cat_id2name, out_csv): rows [] for pred in predictions: file_name pred[file_name].split(/)[-1] for bbox, score, cat_id in zip(pred[instances].pred_boxes.tensor.tolist(), pred[instances].scores.tolist(), pred[instances].pred_classes.tolist()): x, y, w, h bbox rows.append({ image_name: file_name, category_name: cat_id2name[cat_id], score: round(score, 6), xmin: int(x), ymin: int(y), xmax: int(x w), ymax: int(y h), }) pd.DataFrame(rows).to_csv(out_csv, indexFalse)这个文件里最容易错的是浮点精度。比赛平台会对坐标做整数化后计算 IoU如果你直接int()截断而不是round()边界框可能缩水 1~2 像素对APIoU0.5影响不大但对0.5:0.95的影响很直接。建议坐标保留一位小数提交如果不允许就在转换时统一round()后再int()。5.3 后处理的一个进阶技巧按类别自适应置信度阈值Detectron2 默认用的是全局TEST.DETECTIONS_THRESHOLD0.5但比赛数据里不同类别的难易程度差别很大。我习惯在result2to3.py里对不同类别用不同的阈值比如易混淆类别用 0.7简单类用 0.3这样能在不显著增加误检的情况下提升召回。具体做法是先跑一遍验证集统计每个类别的 PR 曲线找到每个类别的最优 score threshold再写进一个 dict 供转换脚本调用。这个技巧常常能比单纯调模型涨 0.5~1 个点而且不需要重新训练。本文还有配套的精品资源点击获取