在 DEIMv2 的目录下创建一个 .py 文件，确保所有组件都成功加载。

import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin

from engine.backbone import HGNetv2, DINOv3STAs
from engine.deim import HybridEncoder, LiteEncoder
from engine.deim import DFINETransformer, DEIMTransformer
from engine.deim.postprocessor import PostProcessor


class DEIMv2(nn.Module, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__()
        self.backbone = DINOv3STAs(**config["DINOv3STAs"])
        self.encoder = HybridEncoder(**config["HybridEncoder"])
        self.decoder = DEIMTransformer(**config["DEIMTransformer"])
        self.postprocessor = PostProcessor(**config["PostProcessor"])

    def forward(self, x, orig_target_sizes):
        x = self.backbone(x)
        x = self.encoder(x)
        x = self.decoder(x)
        x = self.postprocessor(x, orig_target_sizes)

        return x

deimv2_s_config = {
  "DINOv3STAs": {
    ...
  },
  ...
}

deimv2_s_hf = DEIMv2.from_pretrained("Intellindust/DEIMv2_DINOv3_S_COCO")

请按照以下步骤准备 COCO 数据集：

1. 从 OpenDataLab 或 COCO 下载 COCO2017 数据集。

2. 修改 coco_detection.yml 中的路径：

   train_dataloader:
       img_folder: /data/COCO2017/train2017/
       ann_file: /data/COCO2017/annotations/instances_train2017.json
   val_dataloader:
       img_folder: /data/COCO2017/val2017/
       ann_file: /data/COCO2017/annotations/instances_val2017.json
   

若要使用自定义数据集进行训练，需将其整理为 COCO 格式。请按以下步骤准备您的数据集：

1. 将 remap_mscoco_category 设置为 False：

   这将防止类别 ID 自动映射到 MSCOCO 类别。

   remap_mscoco_category: False
   

2. 组织图像文件夹：

   按照以下目录结构整理数据集：

   dataset/
   ├── images/
   │   ├── train/
   │   │   ├── image1.jpg
   │   │   ├── image2.jpg
   │   │   └── ...
   │   ├── val/
   │   │   ├── image1.jpg
   │   │   ├── image2.jpg
   │   │   └── ...
   └── annotations/
       ├── instances_train.json
       ├── instances_val.json
       └── ...
   

   • images/train/：包含所有训练图像。
   • images/val/：包含所有验证图像。
   • annotations/：包含 COCO 格式的标注文件。

3. 将标注转换为 COCO 格式：

   如果您的标注尚未采用 COCO 格式，需要进行转换。您可以参考以下 Python 脚本或使用现有工具：

   import json
   
   def convert_to_coco(input_annotations, output_annotations):
       # 在此处实现转换逻辑
       pass
   
   if __name__ == "__main__":
       convert_to_coco('path/to/your_annotations.json', 'dataset/annotations/instances_train.json')
   

4. 更新配置文件：

   修改您的 custom_detection.yml。

   task: detection
   
   evaluator:
     type: CocoEvaluator
     iou_types: ['bbox', ]
   
   num_classes: 777 # 您的数据集类别数
   remap_mscoco_category: False
   
   train_dataloader:
     type: DataLoader
     dataset:
       type: CocoDetection
       img_folder: /data/yourdataset/train
       ann_file: /data/yourdataset/train/train.json
       return_masks: False
       transforms:
         type: Compose
         ops: ~
     shuffle: True
     num_workers: 4
     drop_last: True
     collate_fn:
       type: BatchImageCollateFunction
   
   val_dataloader:
     type: DataLoader
     dataset:
       type: CocoDetection
       img_folder: /data/yourdataset/val
       ann_file: /data/yourdataset/val/ann.json
       return_masks: False
       transforms:
         type: Compose
         ops: ~
     shuffle: False
     num_workers: 4
     drop_last: False
     collate_fn:
       type: BatchImageCollateFunction
   

1. 训练

   # 对于基于ViT的变体
   CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/deimv2/deimv2_dinov3_${model}_coco.yml --use-amp --seed=0
   
   # 对于基于HGNetv2的变体
   CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/deimv2/deimv2_hgnetv2_${model}_coco.yml --use-amp --seed=0
   

2. 测试

   # 对于基于ViT的变体
   CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/deimv2/deimv2_dinov3_${model}_coco.yml --test-only -r model.pth
   
   # 对于基于HGNetv2的变体
   CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/deimv2/deimv2_hgnetv2_${model}_coco.yml --test-only -r model.pth
   

3. 微调

   # 对于基于ViT的变体
   CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/deimv2/deimv2_dinov3_${model}_coco.yml --use-amp --seed=0 -t model.pth
   
   # 对于基于HGNetv2的变体
   CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --master_port=7777 --nproc_per_node=4 train.py -c configs/deimv2/deimv2_hgnetv2_${model}_coco.yml --use-amp --seed=0 -t model.pth
   

例如，如果您想使用 DEIMv2-S，并在COCO2017上训练 DEIMv2 时将总批量大小增加一倍至64，可以按照以下步骤操作：

1. 修改您的 deimv2_dinov3_s_coco.yml，以增加 total_batch_size：

   train_dataloader:
     total_batch_size: 64 
     dataset: 
       transforms:
         ops:
           ...
   
     collate_fn:
       ...
   

2. 修改您的 deimv2_dinov3_s_coco.yml。以下是关键参数应如何调整：

   optimizer:
     type: AdamW
   
     params: 
       -
         # 除self.dinov3中的归一化、批归一化和偏置外的所有参数
         params: '^(?=.*.dinov3)(?!.*(?:norm|bn|bias)).*$'  
         lr: 0.00005  # 倍增，遵循线性缩放法则
       -
         # 包括self.dinov3中所有归一化、批归一化和偏置
         params: '^(?=.*.dinov3)(?=.*(?:norm|bn|bias)).*$'    
         lr: 0.00005   # 倍增，遵循线性缩放法则
         weight_decay: 0.
       - 
         # 包括除self.dinov3之外的所有归一化、批归一化和偏置
         params: '^(?=.*(?:sta|encoder|decoder))(?=.*(?:norm|bn|bias)).*$'
         weight_decay: 0.
   
     lr: 0.0005   # 如有必要，遵循线性缩放法则
     betas: [0.9, 0.999]
     weight_decay: 0.0001
   
   ema:  # 添加EMA设置
     decay: 0.9998  # 调整为1 - (1 - decay) * 2
     warmups: 500  # 减半
   
   lr_warmup_scheduler:
     warmup_duration: 250  # 减半
   

如果您希望在COCO2017上以320x320的输入尺寸训练 DEIMv2-S，请按照以下步骤操作：

1. 修改您的 deimv2_dinov3_s_coco.yml：

   eval_spatial_size: [320, 320]
   
   train_dataloader:
     # 这里我们以64为例设置总批量大小。
     total_batch_size: 64 
     dataset: 
       transforms:
         ops:
           # 特别是对于马赛克增强，建议输出尺寸为输入尺寸的一半。
           - {type: Mosaic, output_size: 160, rotation_range: 10, translation_range: [0.1, 0.1], scaling_range: [0.5, 1.5],
              probability: 1.0, fill_value: 0, use_cache: True, max_cached_images: 50, random_pop: True}
           ...
           - {type: Resize, size: [320, 320], }
           ...
       collate_fn:
         base_size: 320
         ...
   
   val_dataloader:
     dataset:
       transforms:
         ops:
           - {type: Resize, size: [320, 320], }
           ...
   

如果您想对 DEIMv2-S 进行微调，训练 20 个轮次，请按照以下步骤操作（仅供参考；可根据实际需求进行调整）：

epoches: 32 # 总轮次：20轮训练 + EMA用于4n = 12。n指匹配配置中的模型规模。

flat_epoch: 14    # 4 + 20 // 2
no_aug_epoch: 12  # 4n

train_dataloader:
  dataset: 
    transforms:
      ops:
        ...
      policy:
        epoch: [4, 14, 20]   # [起始轮次、平稳轮次、总轮次 - 无增强轮次]

  collate_fn:
    ...
    mixup_epochs: [4, 14]  # [起始轮次、平稳轮次]
    stop_epoch: 20  # 总轮次 - 无增强轮次
    copyblend_epochs: [4, 20]  # [起始轮次、总轮次 - 无增强轮次]
  
DEIMCriterion:
  matcher:
    ...
    matcher_change_epoch: 18  # 约为(总轮次 - 无增强轮次)的90%

1. 环境准备

   pip install onnx onnxsim
   

2. 导出 ONNX 模型

   python tools/deployment/export_onnx.py --check -c configs/deimv2/deimv2_dinov3_${model}_coco.yml -r model.pth
   

3. 导出 TensorRT 引擎

   trtexec --onnx="model.onnx" --saveEngine="model.engine" --fp16
   

⚠️ TensorRT 版本说明

• ✅ 推荐：使用 TensorRT ≥ 10.6 进行 FP16 推理，以确保运行稳定并获得正确的检测结果。

• ❗ 已知问题：在 TensorRT 10.4 中，FP16 推理可能会产生错误的输出。

• 🔧 旧版本（如 10.4）的解决方法：

  • 以 FP32 模式运行推理，或

  • 仔细验证导出的引擎及端到端流程，以确认数值正确性和检测性能。

1. 环境准备

   pip install -r tools/inference/requirements.txt
   

2. 推理（onnxruntime / tensorrt / torch）

   现已支持对图像和视频进行推理。

   python tools/inference/onnx_inf.py --onnx model.onnx --input image.jpg  # video.mp4
   python tools/inference/trt_inf.py --trt model.engine --input image.jpg
   python tools/inference/torch_inf.py -c configs/deimv2/deimv2_dinov3_${model}_coco.yml -r model.pth --input image.jpg --device cuda:0
   

1. 环境准备

   pip install -r tools/benchmark/requirements.txt
   

2. 模型 FLOPs、MACs 和参数量

   python tools/benchmark/get_info.py -c configs/deimv2/deimv2_dinov3_${model}_coco.yml
   

3. TensorRT 延迟

   python tools/benchmark/trt_benchmark.py --COCO_dir path/to/COCO2017 --engine_dir model.engine
   

1. 环境准备

   pip install fiftyone
   

2. Voxel51 Fiftyone 可视化（fiftyone）

   python tools/visualization/fiftyone_vis.py -c configs/deimv2/deimv2_dinov3_${model}_coco.yml -r model.pth
   

1. 自动恢复训练

   bash reference/safe_training.sh
   

2. 转换模型权重

   python reference/convert_weight.py model.pth