NVIDIA 对象检测工具包 (ODTK)

快速且准确的单阶段对象检测，具备端到端 GPU 优化。

描述

ODTK 是一种单次通过的对象检测器，支持多种骨干网络和检测头。这使得用户可以在性能与精度之间进行权衡。

它针对端到端 GPU 处理进行了优化，使用了以下技术：

• PyTorch 深度学习框架，并支持 ONNX
• NVIDIA Apex，用于混合精度和分布式训练
• NVIDIA DALI，用于优化的数据预处理
• NVIDIA TensorRT，用于高性能推理
• NVIDIA DeepStream，用于优化的实时视频流支持

旋转边界框检测

此仓库现已支持旋转边界框检测。有关如何使用 --rotated-bbox 命令的更多信息，请参阅 旋转检测训练 和 旋转检测推理 文档。

边界框标注由 [x, y, w, h, theta] 描述。

性能

检测流水线允许用户根据所需的延迟-精度权衡选择特定的骨干网络。

ODTK RetinaNet 模型在完成完整训练计划后，针对 COCO 2017（训练/验证集）的精度、推理延迟和 FPS（每秒帧数）。推理结果包括对批量大小为 1 的边界框后处理。推理在 --resize 800 下，使用 --with-dali 并基于 FP16 TensorRT 引擎进行测量。

骨干网络	mAP @[IoU=0.50:0.95]	在 DGX1v 上的训练时间	在 V100 上的 FP16 推理延迟	在 T4 上的 INT8 推理延迟	在 A100 上的 FP16 推理延迟	在 A100 上的 INT8 推理延迟
ResNet18FPN	0.318	5 小时	14 ms； 71 FPS	18 ms； 56 FPS	9 ms； 110 FPS	7 ms； 141 FPS
MobileNetV2FPN	0.333	—	14 ms； 74 FPS	18 ms； 56 FPS	9 ms； 114 FPS	7 ms； 138 FPS
ResNet34FPN	0.343	6 小时	16 ms； 64 FPS	20 ms； 50 FPS	10 ms； 103 FPS	7 ms； 142 FPS
ResNet50FPN	0.358	7 小时	18 ms； 56 FPS	22 ms； 45 FPS	11 ms； 93 FPS	8 ms； 129 FPS
ResNet101FPN	0.376	10 小时	22 ms； 46 FPS	27 ms； 37 FPS	13 ms； 78 FPS	9 ms； 117 FPS
ResNet152FPN	0.393	12 小时	26 ms； 38 FPS	33 ms； 31 FPS	15 ms； 66 FPS	10 ms； 103 FPS

安装

为获得最佳性能，请使用最新的 PyTorch NGC Docker 容器。克隆此仓库，构建并运行您自己的镜像：

git clone https://github.com/nvidia/retinanet-examples
docker build -t odtk:latest retinanet-examples/
docker run --gpus all --rm --ipc=host -it odtk:latest

使用方法

训练、推理、评估和模型导出均可通过 odtk 工具完成。 有关更多详细信息，包括参数列表，请参阅 训练 和 推理 文档。

训练

使用预训练的骨干网络，在 COCO 2017 数据集上训练检测模型：

odtk train retinanet_rn50fpn.pth --backbone ResNet50FPN \
    --images /coco/images/train2017/ --annotations /coco/annotations/instances_train2017.json \
    --val-images /coco/images/val2017/ --val-annotations /coco/annotations/instances_val2017.json

微调

在您的数据集上微调预训练模型。在下面的例子中，我们使用 Pascal VOC 数据集，并采用 JSON 标注：

odtk train model_mydataset.pth --backbone ResNet50FPN \
    --fine-tune retinanet_rn50fpn.pth \
    --classes 20 --iters 10000 --val-iters 1000 --lr 0.0005 \
    --resize 512 --jitter 480 640 --images /voc/JPEGImages/ \
    --annotations /voc/pascal_train2012.json --val-annotations /voc/pascal_val2012.json

注意：输入图像的短边将被调整为 resize 大小，但长边不会超过 max-size。在训练过程中，图像会随机缩放到 jitter 范围内的新尺寸。

推理

在 COCO 2017 数据集上评估您的检测模型：

odtk infer retinanet_rn50fpn.pth --images /coco/images/val2017/ --annotations /coco/annotations/instances_val2017.json

在您的数据集上运行推理：

odtk infer retinanet_rn50fpn.pth --images /dataset/val --output detections.json

使用 TensorRT 进行优化推理

为了加快推理速度，可将检测模型导出为优化后的 FP16 TensorRT 引擎：

odtk export model.pth engine.plan

使用 TensorRT 后端在 COCO 2017 数据集上评估模型：

odtk infer engine.plan --images /coco/images/val2017/ --annotations /coco/annotations/instances_val2017.json

使用 TensorRT 进行 INT8 推理

若需更快速的推理，可进行 INT8 校准以创建优化后的 INT8 TensorRT 引擎：

odtk export model.pth engine.plan --int8 --calibration-images /coco/images/val2017/

这将生成一个 INT8CalibrationTable 文件，以后可用于为同一模型创建 INT8 TensorRT 引擎，而无需再次校准。

或者，使用缓存的校准表创建优化后的 INT8 TensorRT 引擎：

odtk export model.pth engine.plan --int8 --calibration-table /path/to/INT8CalibrationTable

数据集

RetinaNet 支持 COCO JSON 格式的标注。将您自己的数据集标注转换为 JSON 格式时，需要包含以下字段：

{
    "images": [{
        "id" : int,
        "file_name" : str
    }],
    "annotations": [{
        "id" : int,
        "image_id" : int, 
        "category_id" : int,
        "bbox" : [x, y, w, h]   # 全部为浮点数
        "area": float           # w * h。用于计算验证分数时必需
        "iscrowd": 0            # 用于计算验证分数时必需
    }],
    "categories": [{
        "id" : int
    ]}
}

如果使用 --rotated-bbox 标志进行旋转目标检测，则需要在标注中添加额外的浮点数 theta。若要获得验证分数，还需填写 segmentation 部分。

        "bbox" : [x, y, w, h, theta]    # 全部为浮点数，其中 theta 以弧度表示，从 x 轴逆时针方向测量。
        "segmentation" : [[x1, y1, x2, y2, x3, y3, x4, y4]]
                                        # 用于计算验证分数时必需。

免责声明

这是一个研究项目，并非 NVIDIA 的官方产品。

Jetpack 兼容性

本分支使用 TensorRT 7。如果您使用 PyTorch 训练和推理模型，或者在 Tesla GPU（例如 V100、T4）上创建 TensorRT 引擎，则应使用此分支。

如果您希望将模型部署到运行 Jetpack 4.3 版本的 Jetson 设备（例如 Jetson AGX Xavier），则应使用本仓库的 19.10 分支。

参考文献

• 密集目标检测中的焦点损失。 林宗毅、普里亚·戈亚尔、罗斯·吉尔希克、何凯明、皮奥特·多拉尔。 ICCV，2017 年。
• 准确的大批量 SGD：1 小时内训练 ImageNet。 普里亚·戈亚尔、皮奥特·多拉尔、罗斯·吉尔希克、彼得·诺德胡伊斯、卢卡什·韦索洛夫斯基、阿波·基罗拉、安德鲁·图洛克、贾扬青、何凯明。 2017 年 6 月。
• 用于目标检测的特征金字塔网络。 林宗毅、皮奥特·多拉尔、罗斯·吉尔希克、何凯明、巴拉特·哈里哈兰、塞尔日·贝隆吉。 CVPR，2017 年。
• 用于图像识别的深度残差学习。 何凯明、张祥宇、任少卿、孙健。 CVPR，2016 年。

RetinaNet Examples (NVIDIA ODTK) 快速上手指南

本指南基于 NVIDIA 目标检测工具包 (ODTK)，帮助开发者快速部署和使用高性能的 RetinaNet 模型进行单阶段物体检测。

环境准备

为了获得最佳性能，强烈建议使用 NVIDIA 提供的 PyTorch NGC Docker 容器。该容器已预装 PyTorch、ONNX、Apex、DALI 和 TensorRT 等必要依赖。

• 系统要求：
  • Linux 操作系统
  • NVIDIA GPU (支持 CUDA)
  • 已安装 NVIDIA Docker (nvidia-docker2 或 Docker 支持 --gpus 参数)
• 前置依赖：
  • Docker Engine
  • NVIDIA Container Toolkit

注意：本项目为研究性质代码，非 NVIDIA 官方正式产品。若需部署至 Jetson 设备 (如 Jetson AGX Xavier) 且运行 Jetpack 4.3，请使用 19.10 分支；其他 Tesla GPU (V100, T4, A100) 训练及推理请使用主分支 (基于 TensorRT 7)。

安装步骤

通过克隆仓库并构建 Docker 镜像即可完成安装。国内用户若拉取 NGC 镜像较慢，可先配置 Docker 国内加速源。

1. 克隆代码仓库：

   git clone https://github.com/nvidia/retinanet-examples
   

2. 构建 Docker 镜像：

   docker build -t odtk:latest retinanet-examples/
   

3. 启动容器：

   docker run --gpus all --rm --ipc=host -it odtk:latest
   

基本使用

进入容器后，所有操作（训练、推理、评估、模型导出）均通过 odtk 命令行工具完成。数据标注需遵循 COCO JSON 格式。

1. 模型训练 (Training)

使用预训练的骨干网络在 COCO 2017 数据集上训练检测模型：

odtk train retinanet_rn50fpn.pth --backbone ResNet50FPN \
    --images /coco/images/train2017/ --annotations /coco/annotations/instances_train2017.json \
    --val-images /coco/images/val2017/ --val-annotations /coco/annotations/instances_val2017.json

2. 模型微调 (Fine Tuning)

在自定义数据集（如下方的 Pascal VOC 示例）上微调预训练模型：

odtk train model_mydataset.pth --backbone ResNet50FPN \
    --fine-tune retinanet_rn50fpn.pth \
    --classes 20 --iters 10000 --val-iters 1000 --lr 0.0005 \
    --resize 512 --jitter 480 640 --images /voc/JPEGImages/ \
    --annotations /voc/pascal_train2012.json --val-annotations /voc/pascal_val2012.json

3. 模型推理 (Inference)

对验证集或自定义图片目录进行推理：

# 评估 COCO 验证集
odtk infer retinanet_rn50fpn.pth --images /coco/images/val2017/ --annotations /coco/annotations/instances_val2017.json

# 推理自定义数据集并输出结果
odtk infer retinanet_rn50fpn.pth --images /dataset/val --output detections.json

4. 高性能推理 (TensorRT 加速)

将模型导出为优化的 TensorRT 引擎以大幅提升推理速度：

导出 FP16 引擎：

odtk export model.pth engine.plan

导出 INT8 引擎 (需校准)：

# 创建校准表并生成引擎
odtk export model.pth engine.plan --int8 --calibration-images /coco/images/val2017/

# 或使用已有的校准表生成引擎
odtk export model.pth engine.plan --int8 --calibration-table /path/to/INT8CalibrationTable

使用 TensorRT 引擎进行推理：

odtk infer engine.plan --images /coco/images/val2017/ --annotations /coco/annotations/instances_val2017.json

提示：如需支持旋转框检测，请在训练和推理命令中添加 --rotated-bbox 参数，并确保标注格式包含 [x, y, w, h, theta]。