ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

DiT 文档布局分析实战:基于 Detectron2 的 Mask R-CNN / Cascade Mask R-CNN 推理、训练与评估完全指南

DiT 文档布局分析实战:基于 Detectron2 的 Mask R-CNN / Cascade Mask R-CNN 推理、训练与评估完全指南 DiT 文档布局分析实战基于 Detectron2 的 Mask R-CNN / Cascade Mask R-CNN 推理、训练与评估完全指南【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm本指南以 unilm 仓库中 dit/object_detection 模块为核心系统讲解如何使用 DiTDocument Image Transformer作为主干网络在 Detectron2 框架上完成文档布局分析Document Layout Analysis中的目标检测任务覆盖 PubLayNet 与 ICDAR 2019 cTDaR 两大数据集的推理、数据准备、评估与微调全流程。读者读完本文后将能够独立复现 DiT-Base / DiT-Large 搭配 Mask R-CNN / Cascade Mask R-CNN 的检测方案并理解其配置体系与底层实现原理。1. 模块概述DiT 的文档布局分析实现dit/object_detection是 DiTDocument Image Transformer在文档布局分析任务上的官方实现。该目录基于 Meta 的 Detectron2Mask R-CNN 与 Cascade Mask R-CNN 的实现来源将 DiT 预训练权重作为视觉主干backbone接入两阶段检测框架面向两个文档数据集PubLayNet大规模文档版面数据集检测类别为 5 类 ——text、title、list、table、figureICDAR 2019 cTDaR表格检测与识别竞赛数据集目标类别为table区分 modern 与 archival 两个子集。目录的核心文件布局如下路径作用inference.py单图推理与结果可视化脚本train_net.py训练 / 评估统一入口convert_to_coco_format.py将 ICDAR 2019 cTDaR 原始标注转为 COCO 格式adaptive_binarize.py对 archival 子集做自适应二值化publaynet_configsPubLayNet 的 Mask R-CNN / Cascade Mask R-CNN 配置icdar19_configsICDAR 2019 cTDaR 的 Mask R-CNN / Cascade Mask R-CNN 配置ditodDiT 主干、数据集映射、评估器与训练器扩展实现publaynet_example.jpeg推理演示示例图片其中ditod子包是整个方案的“引擎舱”包含 backbone.pyViT FPN 主干、beit.py 与 deit.pyDiT / BEiT / DEiT / MAE 模型定义、config.pyViT 专属配置项、dataset_mapper.pyDETR 式数据增强、mytrainer.py自定义训练器以及 icdar_evaluation.pyICDAR 评估器。2. 推理快速体验 DiT 文档布局分析2.1 使用 inference.py 进行单图推理官方提供了 Hugging Face Spaces 网页演示可直接在线体验文档布局分析效果而在本地最快的验证方式是运行inference.py脚本。以下命令需在unilm 仓库根目录执行python ./dit/object_detection/inference.py \ --image_path ./dit/object_detection/publaynet_example.jpeg \ --output_file_name output.jpg \ --config ./dit/object_detection/publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml \ --opts MODEL.WEIGHTS https://layoutlm.blob.core.windows.net/dit/dit-fts/publaynet_dit-b_mrcnn.pth务必保证配置YAML与 PyTorch 权重匹配。上例使用的是 DiT-Base 主干 Mask R-CNN 框架在 PubLayNet 上微调后的权重publaynet_dit-b_mrcnn.pth若改用 DiT-Large 配置则需替换为对应的dit_large权重否则加载 checkpoint 时会因网络结构不匹配而失败。四个命令行参数的含义分别为参数说明--image_path输入图片路径必填--output_file_name可视化结果输出文件名如output.jpg--config检测框架配置文件路径--opts覆盖配置项格式为KEY VALUE键值对此处用于指定微调权重2.2 推理脚本源码解析从 inference.py 的源码可以看到推理的完整链路构造配置get_cfg()创建 Detectron2 默认配置随后调用add_vit_config(cfg)注入MODEL.VIT.*等 DiT 专属配置项见 config.py再merge_from_file读取 YAML、merge_from_list应用--opts覆盖设定设备device cuda if torch.cuda.is_available() else cpu自动回退到 CPUCPU 推理速度较慢仅适合快速验证构建预测器DefaultPredictor(cfg)由 Detectron2 提供会自动加载MODEL.WEIGHTS指定权重并对输入做ResizeShortestEdge预处理设置类别元数据根据测试数据集名动态指定类别 —— 若cfg.DATASETS.TEST[0] icdar2019_test则类别为[table]否则为[text,title,list,table,figure]PubLayNet 五类推理与可视化Visualizer以ColorMode.SEGMENTATION模式绘制预测实例框 掩码 类别最终通过cv2.imwrite保存。这一流程清晰展示了 Detectron2 “配置驱动”的工程范式模型结构完全由配置文件决定脚本只负责组装。3. 数据集准备3.1 PubLayNetPubLayNet 数据集约 96GB需从官方发布渠道下载publaynet.tar.gz后解压到目录PATH-to-PubLayNet。解压后执行ln -s PATH-to-PubLayNet publaynet_data软链接名称必须为publaynet_data。其原因在 train_net.py 中写死脚本通过register_coco_instances注册数据集时硬编码了./publaynet_data/train.json、./publaynet_data/train等相对路径。因此在 unilm 仓库根目录下创建该软链接程序才能访问到数据。3.2 ICDAR 2019 cTDaRICDAR 2019 cTDaR 数据约 4GB下载后假设仓库路径名为PATH-to-ICDARrepo。首先将原始数据转换为 COCO 格式python convert_to_coco_format.py --root_dirPATH-to-ICDARrepo --target_dirPATH-toICDAR处理后的数据位于PATH-to-ICDAR。接着对archival 子集执行自适应二值化现代印刷体 modern 子集无需处理cp -r PATH-to-ICDAR/trackA_archival PATH-to-ICDAR/at_trackA_archival python adaptive_binarize.py --root_dir PATH-to-ICDAR/at_trackA_archival二值化后的 archival 子集保存在PATH-to-ICDAR/at_trackA_archival。随后根据要评估/微调的子集在仓库根目录建立data软链接ln -s PATH-to-ICDAR/trackA_modern data # 评估 modern 子集 # 或 ln -s PATH-to-ICDAR/at_trackA_archival data # 评估 archival 子集与 PubLayNet 同理train_net.py 中注册 ICDAR 数据时使用的是data/train.json、data/test.json等相对路径因此软链接必须命名为data且建立在当前工作目录。3.3 数据预处理脚本源码解读convert_to_coco_format.py 的核心逻辑是将 ICDAR 的 XML 标注解析为 COCO JSON类别固定为单一tablecategories: [{id: 1, name: table}]从 XML 的table/Coords节点读取表格四角点计算 segmentation 多边形与bbox取[x0, y0, x3-x0, y3-y0]文件名前缀cTDaR_t0对应trackA_archival、cTDaR_t1对应trackA_modern。该脚本还内置clean_img()函数用于统一.JPG、.TIFF、.png等图片格式为.jpg。adaptive_binarize.py 使用 OpenCV 的cv2.adaptiveThresholdADAPTIVE_THRESH_GAUSSIAN_CblockSize45C11对灰度图做高斯自适应阈值二值化再转回三通道 BGR 覆写原图以提升档案扫描件的表格边界检测效果。4. 评估验证微调后的检测性能评估使用 train_net.py 的--eval-only模式。配置文件位于icdar19_configs与publaynet_configs两个目录。示例 1评估 PubLayNet 上微调的 DiT-Base Mask R-CNNpython train_net.py --config-file publaynet_configs/maskrcnn/maskrcnn_dit_base.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS finetuned_checkpoint_file_path or link OUTPUT_DIR your_output_dir示例 2评估 ICDAR 2019 cTDaR archival 子集上微调的 DiT-Large Cascade Mask R-CNN需先将PATH-to-ICDAR/at_trackA_archival软链接为datapython train_net.py --config-file icdar19_configs/cascade/cascade_dit_large.yaml --eval-only --num-gpus 8 MODEL.WEIGHTS finetuned_checkpoint_file_path or link OUTPUT_DIR your_output_dir4.1 ICDAR 2019 测量工具的 Bug 修复重要提示官方在将 ICDAR2019 测量工具集成进代码时修复了原工具中的一个 bug。如果你使用外部 ICDAR 测量工具ctdar_measurement_tool自行计算评估分数请将evaluate.py中按扩展名过滤 ground-truth 文件的代码修改如下原代码for file in gt_file_lst: ... gt_file_lst.remove(file)在遍历过程中删除列表元素会导致漏删或越界... # print(each_file) # for file in gt_file_lst: # if file.split(.) ! xml: # gt_file_lst.remove(file) # # print(gt_file_lst) # Comment the code above and add the code below for i in range(len(gt_file_lst) - 1, -1, -1): if gt_file_lst[i].split(.)[-1] ! xml: del gt_file_lst[i] if len(gt_file_lst) 0: ...即改为从后往前倒序遍历并删除避免“边遍历边删除”造成的元素遗漏问题。4.2 评估器实现仓库自带的评估分流逻辑位于 mytrainer.py 的build_evaluator数据集名包含icdar时使用自定义ICDAREvaluator见 icdar_evaluation.py其内部集成了修复后的 ICDAR 测量逻辑其余情况使用 Detectron2 标准COCOEvaluator。5. 训练微调 DiT 主干以下两条命令展示了如何使用 DiT 主干 Mask R-CNN / Cascade Mask R-CNN 在8 张 32GB NVIDIA V100 GPU上进行微调。示例 1PubLayNet 上微调 DiT-Base Cascade Mask R-CNNpython train_net.py --config-file publaynet_configs/cascade/cascade_dit_base.yaml --num-gpus 8 MODEL.WEIGHTS DiT-Base_file_path or link OUTPUT_DIR your_output_dir示例 2ICDAR 2019 cTDaR modern 子集上微调 DiT-Large Mask R-CNNpython train_net.py --config-file icdar19_configs/markrcnn/maskrcnn_dit_large.yaml --num-gpus 8 MODEL.WEIGHTS DiT-Large_file_path or link OUTPUT_DIR your_output_dir微调时MODEL.WEIGHTS传入的是DiT 自监督预训练权重如dit-base-224-p16-500k-62d53a.pth、dit-large-224-p16-500k-d7a2fb.pth由 maskrcnn_dit_base.yaml 等配置的MODEL.WEIGHTS字段指定命令行的MODEL.WEIGHTS覆盖则用于指定已微调 checkpoint配合--eval-only或替换预训练权重来源。更详细的 Detectron2 用法可参考其官方文档。5.1 训练入口与数据集注册train_net.py 的main()首先通过register_coco_instances注册四个数据集publaynet_train/val、icdar2019_train/test随后setup(args)完成配置合并与冻结。它复用了 Detectron2 的launch()分布式启动器支持--num-gpus、--num-machines、--machine-rank、--dist-url等标准参数并额外提供--debug参数内部使用 debugpy 在 0.0.0.0:9310 等待调试器附加。5.2 自定义训练器与优化策略MyTrainer 继承自 Detectron2 的TrainerBase其中几个关键设计点数据加载当cfg.AUG.DETRTrue时使用自定义 DetrDatasetMapper启用 DETR 风格增强 —— 以 50% 概率插入ResizeShortestEdge([400,500,600]) RandomCrop(absolute_range, (384,600))裁剪序列AMP 混合精度cfg.SOLVER.AMP.ENABLEDTrue时训练循环自动切换为AMPTrainer优化器支持 SGD / AdamW并为 backbone 参数提供BACKBONE_MULTIPLIER学习率缩放启用full_model梯度裁剪时会在 step 前对整个模型参数执行clip_grad_norm_见build_optimizer调度器使用WarmupCosineLR配置了WARMUP_FACTOR0.01、WARMUP_ITERS等钩子内置IterationTimer、LRScheduler、PreciseBN、PeriodicCheckpointer、EvalHook、PeriodicWriter等训练钩子并在训练期间按TEST.EVAL_PERIOD自动做周期性评估。6. 配置文件逐项解析配置文件采用 Detectron2 的 YACS 继承体系子配置通过_BASE_: ../Base-RCNN-FPN.yaml继承公共配置再按数据集与模型规模覆盖差异项。6.1 公共配置 Base-RCNN-FPN.yamlpublaynet_configs/Base-RCNN-FPN.yaml 定义了检测框架的公共结构配置项值说明MODEL.META_ARCHITECTUREGeneralizedRCNN标准两阶段检测架构MODEL.MASK_ONTrue启用实例分割分支MODEL.PIXEL_MEAN/STD[123.675, 116.280, 103.530]/[58.395, 57.120, 57.375]图像归一化参数DiT 配置会覆盖为 127.5 系MODEL.BACKBONE.NAMEbuild_vit_fpn_backbone注册到BACKBONE_REGISTRY的 ViTFPN 主干MODEL.VIT.OUT_FEATURES[layer3,layer5,layer7,layer11]从 DiT 提取的多尺度特征层MODEL.VIT.DROP_PATH0.1随机深度Stochastic Depth丢弃率MODEL.VIT.IMG_SIZE[224,224]预训练输入分辨率MODEL.VIT.POS_TYPEabs绝对位置编码MODEL.FPN.IN_FEATURES与OUT_FEATURES相同FPN 输入特征MODEL.ROI_HEADS.NUM_CLASSES5PubLayNet 五类SOLVER.BASE_LR0.0004基础学习率SOLVER.IMS_PER_BATCH32全局 batch sizeINPUT.CROPabsolute_range (384,600)DETR 式随机裁剪INPUT.MIN_SIZE_TRAIN(480,512,...,800)短边随机缩放范围AUG.DETRTrue启用 DETR 数据增强SEED42随机种子6.2 VIT 配置注入add_vit_config所有配置文件都必须先经 config.py 的add_vit_config(cfg)注入MODEL.VIT.*默认值否则会出现“配置项不存在”错误。其注册的默认值包括MODEL.VIT.NAME默认主干模型名可选dit_base_patch16、dit_large_patch16、beit_base_patch16、beit_large_patch16、deit_base_patch16、mae_base_patch16MODEL.VIT.OUT_FEATURES默认[layer3,layer5,layer7,layer11]输出哪些 Transformer 层的特征MODEL.VIT.IMG_SIZE默认[224,224]MODEL.VIT.POS_TYPE默认shared_rel位置编码类型可取值abs/shared_rel/relMODEL.VIT.DROP_PATH默认0.MODEL.VIT.MODEL_KWARGS默认{}透传给模型构造函数的额外参数SOLVER.OPTIMIZER默认ADAMW、SOLVER.BACKBONE_MULTIPLIER默认1.0AUG.DETR默认False是否启用 DETR 数据增强。6.3 各数据集与模型规模的配置差异PubLayNet Mask R-CNNDiT-Base—— maskrcnn_dit_base.yaml覆盖PIXEL_MEAN/STD为[127.5, 127.5, 127.5]与 DiT 预训练归一化一致MODEL.VIT.NAME: dit_base_patch16WARMUP_ITERS: 1000、IMS_PER_BATCH: 16、MAX_ITER: 60000、CHECKPOINT_PERIOD: 2000、TEST.EVAL_PERIOD: 2000。PubLayNet Cascade Mask R-CNNDiT-Base—— cascade_dit_base.yaml在 Mask R-CNN 基础上将ROI_HEADS.NAME改为CascadeROIHeads、ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG: True类别无关的框回归、RPN.POST_NMS_TOPK_TRAIN: 2000。ICDAR 2019DiT-Large—— maskrcnn_dit_large.yaml 与 cascade_dit_large.yaml 的共同差异MODEL.VIT.NAME: dit_large_patch16OUT_FEATURES与FPN.IN_FEATURES切换为[layer7,layer11,layer15,layer23]24 层 DiT-Large 的深层特征DROP_PATH: 0.2学习率降至BASE_LR: 0.00005IMS_PER_BATCH: 16checkpoint / 评估周期缩短为1000。7. DiT 主干网络实现原理7.1 VIT_Backbone 与 FPN 的组装backbone.py 中的VIT_Backbone负责将 ViT 模型包装为 Detectron2 的Backbone其_out_feature_strides按模型规模区分Base 系列dit_base_patch16等12 层layer3→stride 4、layer5→8、layer7→16、layer11→32Large 系列dit_large_patch16、beit_large_patch1624 层layer7→4、layer11→8、layer15→16、layer23→32。build_vit_fpn_backbone注册为build_vit_fpn_backbone在VIT_Backbone之上叠加 Detectron2 标准FPNtop_block使用LastLevelMaxPool生成p6层最终形成 P2–P6 特征金字塔供 RPN 与 ROI Heads 使用。7.2 DiT 模型结构与多尺度特征输出DiT 的模型定义位于 beit.py 的BEiT类。从源码看dit_base_patch16与dit_large_patch16的差异主要体现在embed_dim / 深度 / 头数Base 为 768 / 12 层 / 12 头Large 为 1024 / 24 层 / 16 头LayerScale 初值Base 为init_values0.1Large 为init_values1e-5残差分支乘以可学习缩放向量见Block中的gamma_1/gamma_2两者均使用qkv_biasTrue、patch_size16、mlp_ratio4。forward_features逐层前向当层号命中out_indices时将 token 序列重排回二维特征图去掉cls_token后reshape(B, C, Hp, Wp)最后经过四个轻量 FPN 头生成多尺度输出patch16 场景下fpn1为两层ConvTranspose2d上采样stride 4、fpn2为单层上采样stride 8、fpn3为Identitystride 16、fpn4为MaxPool2dstride 32对应 backbone 中声明的主干 stride 映射。use_checkpointTrue时各 Block 通过torch.utils.checkpoint做激活重计算以节省显存这也是 8×V100 32GB 能跑 DiT-Large 的重要原因。7.3 位置编码与推理灵活性BEiT支持三种位置编码对应配置POS_TYPEabs可学习的绝对位置编码use_abs_pos_embTrueshared_rel跨层共享的相对位置偏置RelativePositionBias在Attention中加到注意力分数上rel每层独立的窗口内相对位置偏置。其中RelativePositionBias实现了 bicubic 插值当推理分辨率与预训练[224,224]不一致时可自动将位置偏置表插值到新的窗口尺寸见 beit.py 中training_window_size ! window_size的分支。配合 deit.py 中的interpolate_pos_encoding使得 DiT 主干能处理文档检测所需的任意分辨率输入。8. 注意事项与常见问题配置与权重必须匹配DiT-Base 配置配 DiT-Base 预训练/微调权重DiT-Large 同理混用会导致 checkpoint 加载失败或精度异常。数据集软链接命名必须在运行命令的目录下创建publaynet_dataPubLayNet与dataICDAR软链接因为 train_net.py 中数据集注册路径是硬编码的相对路径。archival 子集需要二值化仅对 ICDAR 的trackA_archival执行adaptive_binarize.pymodern 子集直接使用原始扫描图。ICDAR 测量工具 Bug若使用第三方 ctdar_measurement_tool 复算分数必须按上文修复evaluate.py中 gt 文件过滤逻辑。运行环境本模块依赖 Detectron2Mask R-CNN / Cascade Mask R-CNN 实现与 timm 库训练推荐 8×32GB V100配置中的IMS_PER_BATCH、BASE_LR、MAX_ITER以 8 卡为基准。归一化参数DiT 配置将PIXEL_MEAN/STD覆盖为[127.5, 127.5, 127.5]与 DiT 预训练mean0.5、std0.5等价于归一化到 [-1,1]保持一致切勿沿用 Detectron2 默认的 ImageNet 归一化。9. 引用与致谢如果本仓库对您的研究或工程有所帮助请引用 DiT 论文misc{li2022dit, title{DiT: Self-supervised Pre-training for Document Image Transformer}, author{Junlong Li and Yiheng Xu and Tengchao Lv and Lei Cui and Cha Zhang and Furu Wei}, year{2022}, eprint{2203.02378}, archivePrefix{arXiv}, primaryClass{cs.CV} }特别感谢 Detectron2 项目提供的 Mask R-CNN 与 Cascade Mask R-CNN 实现以及 DETR / DINO / timm / BEiT 等开源工作为 DiT 检测分支带来的工程基础。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进