YolactEdge：边缘设备上的实时实例分割

██╗   ██╗ ██████╗ ██╗      █████╗  ██████╗████████╗    ███████╗██████╗  ██████╗ ███████╗
╚██╗ ██╔╝██╔═══██╗██║     ██╔══██╗██╔════╝╚══██╔══╝    ██╔════╝██╔══██╗██╔════╝ ██╔════╝
 ╚████╔╝ ██║   ██║██║     ███████║██║        ██║       █████╗  ██║  ██║██║  ███╗█████╗  
  ╚██╔╝  ██║   ██║██║     ██╔══██║██║        ██║       ██╔══╝  ██║  ██║██║   ██║██╔══╝  
   ██║   ╚██████╔╝███████╗██║  ██║╚██████╗   ██║       ███████╗██████╔╝╚██████╔╝███████╗
   ╚═╝    ╚═════╝ ╚══════╝╚═╝  ╚═╝ ╚═════╝   ╚═╝       ╚══════╝╚═════╝  ╚═════╝ ╚══════╝

YolactEdge 是首个能够在小型边缘设备上以实时速度运行的具有竞争力的实例分割方法。具体而言，YolactEdge 在 Jetson AGX Xavier 上最高可达到 30.8 FPS（在 RTX 2080 Ti 上则为 172.7 FPS），使用 ResNet-101 作为骨干网络，处理分辨率为 550x550 的图像。这是 我们的论文 对应的代码。

如需实时演示及更多示例，请观看我们的演示视频。

模型库

我们提供了基于 COCO 和 YouTube VIS 数据集（我们的子训练集，并与 COCO 联合训练）训练的 YOLACT 和 YolactEdge 基线模型。

要评估模型，请将相应的权重文件放入 ./weights 目录，并运行以下命令之一。

YouTube VIS 模型：

方法	骨干网络	mAP	AGX-Xavier FPS	RTX 2080 Ti FPS	权重
YOLACT	R-50-FPN	44.7	8.5	59.8	下载 | 镜像
YolactEdge (不使用 TRT)	R-50-FPN	44.2	10.5	67.0	下载 | 镜像
YolactEdge	R-50-FPN	44.0	32.4	177.6	下载 | 镜像
YOLACT	R-101-FPN	47.3	5.9	42.6	下载 | 镜像
YolactEdge (不使用 TRT)	R-101-FPN	46.9	9.5	61.2	下载 | 镜像
YolactEdge	R-101-FPN	46.2	30.8	172.7	下载 | 镜像

COCO 模型：

方法	骨干网络	mAP	Titan Xp FPS	AGX-Xavier FPS	RTX 2080 Ti FPS	权重
YOLACT	MobileNet-V2	22.1	-	15.0	35.7	下载 | 镜像
YolactEdge	MobileNet-V2	20.8	-	35.7	161.4	下载 | 镜像
YOLACT	R-50-FPN	28.2	42.5	9.1	45.0	下载 | 镜像
YolactEdge	R-50-FPN	27.0	-	30.7	140.3	下载 | 镜像
YOLACT	R-101-FPN	29.8	33.5	6.6	36.5	下载 | 镜像
YolactEdge	R-101-FPN	29.5	-	27.3	124.8	下载 | 镜像

安装

请参阅 INSTALL.md。

您也可以选择使用官方提供的 Dockerfile，通过一条命令即可搭建完整的环境。

快速入门

请按照 安装说明 设置运行 YolactEdge 所需的环境。

请参阅关于如何使用 YolactEdge 进行 评估 和 训练 的说明。

Colab 笔记本

您可以尝试我们的 Colab 笔记本，其中包含实时演示，帮助您了解基本用法。

如果您对使用 TensorRT 评估 YolactEdge 感兴趣，我们还提供另一份 Colab 笔记本，其中配置了 Colab 上的 TensorRT 环境。

评估

定量结果

# 使用最佳设置将训练好的模型的各个组件转换为 TensorRT，并在 YouTube VIS 验证集（我们的划分）上进行评估。
python3 eval.py --trained_model=./weights/yolact_edge_vid_847_50000.pth

# 在整个 COCO 验证集上进行评估。
python3 eval.py --trained_model=./weights/yolact_edge_54_800000.pth

# 输出用于 COCO test-dev 的 COCO JSON 文件。该命令将分别生成用于检测和实例分割的 './results/bbox_detections.json' 和 './results/mask_detections.json' 文件。这些文件随后可以提交到网站进行评估。
python3 eval.py --trained_model=./weights/yolact_edge_54_800000.pth --dataset=coco2017_testdev_dataset --output_coco_json

定性结果

# 在 COCO 数据集上展示定性结果。从这里开始我将使用 0.3 的置信度阈值。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --display

基准测试

# 在 COCO 验证集上对训练好的模型进行基准测试。
# 只运行原始模型，在验证集的前 1000 张图像上执行。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --benchmark --max_images=1000

注释

使用 TensorRT 时的推理错误处理

如果您在使用 YolactEdge 的 TensorRT 转换时，在后处理或 NMS 阶段遇到问题，这可能与 TensorRT 引擎本身的问题有关。我们实现了一个实验性的安全模式，可以谨慎地处理这些情况。请在命令中使用 --use_tensorrt_safe_mode 选项尝试一下。

使用 YOLACT 训练的模型进行推理

如果您有一个使用 YOLACT 训练的预训练模型，并希望利用 YolactEdge 的 TensorRT 功能，只需在命令行选项中指定 --config=yolact_edge_config，代码将自动检测并转换模型权重以确保兼容性。

python3 eval.py --config=yolact_edge_config --trained_model=./weights/yolact_base_54_800000.pth

不进行校准的推理

如果您希望在不进行校准的情况下运行推理命令，可以选择仅使用 FP16 的 TensorRT 优化，或者使用相应的配置而不启用 TensorRT 优化。有关此类配置的示例，请参阅 data/config.py。

# 使用仅 FP16 的 TensorRT 优化评估 YolactEdge，通过 `--use_fp16_tensorrt` 选项（将所有 INT8 优化替换为 FP16）。
python3 eval.py --use_fp16_tensorrt --trained_model=./weights/yolact_edge_54_800000.pth

# 使用 `--disable_tensorrt` 选项评估不使用 TensorRT 优化的 YolactEdge。
python3 eval.py --disable_tensorrt --trained_model=./weights/yolact_edge_54_800000.pth

图像

# 在指定图像上显示定性结果。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --image=my_image.png

# 处理一张图像并将其保存到另一个文件。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --image=input_image.png:output_image.png

# 处理整个文件夹中的图像。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --images=path/to/input/folder:path/to/output/folder

视频

# 实时显示视频。`--video_multiframe` 将同时处理多个帧以提高性能。
# 如果 video_multiframe > 1，则应相应增加 trt_batch_size 以匹配或超过该值。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --video_multiframe=2 --trt_batch_size 2 --video=my_video.mp4

# 实时显示网络摄像头画面。如果有多个摄像头，请指定您想要使用的摄像头索引，而不是 0。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --video_multiframe=2 --trt_batch_size 2 --video=0

# 处理视频并将其保存到另一个文件。此操作未经过优化。
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --video=input_video.mp4:output_video.mp4

使用帮助选项可查看所有可用命令行参数的说明：

python eval.py --help

程序化推理

您可以将 yolact_edge 作为包集成到自己的代码中。为此需要两个步骤：

1. 将 YOLACT Edge 安装为 Python 包：pip install .
2. 按照 pkg_usage.py 中提供的示例使用它。

训练

请确保使用上述命令下载完整的数据集。

• 要进行训练，需获取一个 ImageNet 预训练模型并将其放入 ./weights 目录。
  • 对于 Resnet101，从 这里 下载 resnet101_reducedfc.pth。
  • 对于 Resnet50，从 这里 下载 resnet50-19c8e357.pth。
  • 对于 MobileNetV2，从 这里 下载 mobilenet_v2-b0353104.pth。
• 运行以下任一训练命令。
  • 请注意，在训练过程中按 Ctrl+C 可中断训练，系统会保存当前迭代的 *_interrupt.pth 文件。
  • 默认情况下，所有权重都会保存在 ./weights 目录中，文件名为 <config>_<epoch>_<iter>.pth。

# 使用基础边缘配置，批次大小为 8（默认值）进行训练。
python train.py --config=yolact_edge_config

# 使用特定权重文件恢复训练，并从权重文件名中指定的迭代开始继续训练。
python train.py --config=yolact_edge_config --resume=weights/yolact_edge_10_32100.pth --start_iter=-1

# 使用帮助选项查看所有可用命令行参数的说明。
python train.py --help

基于视频数据集的训练

# 对基于图像的模型进行预训练。
python train.py --config=yolact_edge_youtubevis_config

# 训练流（变形）模块。
python train.py --config=yolact_edge_vid_trainflow_config --resume=./weights/yolact_edge_youtubevis_847_50000.pth

# 联合微调网络。
python train.py --config=yolact_edge_vid_config --resume=./weights/yolact_edge_vid_trainflow_144_100000.pth

自定义数据集

你也可以按照以下步骤使用自己的数据集进行训练：

• 根据你的数据集类型，为你的数据集创建 COCO 样式（图像）或 YTVIS 样式（视频）的目标检测 JSON 注释文件。相关规范分别可在 COCO 和 YTVIS 的文档中找到。请注意，我们不会使用某些字段，因此可以省略以下内容：
  • info
  • license
  • 在 image 下：license, flickr_url, coco_url, date_captured
  • categories（我们使用自定义的类别格式，见下文）
• 在 data/config.py 文件中的 dataset_base 下为你的数据集创建定义（请参阅 dataset_base 中的注释以了解每个字段的说明）：

my_custom_dataset = dataset_base.copy({
    'name': '我的数据集',

    'train_images': '训练图像路径',
    'train_info':   '训练标注文件路径',

    'valid_images': '验证图像路径',
    'valid_info':   '验证标注文件路径',

    'has_gt': True,
    'class_names': ('my_class_id_1', 'my_class_id_2', 'my_class_id_3', ...),

    # 以下仅适用于 YTVIS 样式的视频数据集。

    # 是否采样所有帧或仅关键帧。
    'use_all_frames': False,

    # 以下四行定义了给定数据集的帧采样策略。
    'frame_offset_lb': 1,
    'frame_offset_ub': 4,
    'frame_offset_multiplier': 1,
    'all_frame_direction': 'allway',

    # 每段视频标注 K 分之 1 帧
    'images_per_video': 5,

    # 声明这是一个视频数据集
    'is_video': True
})

• 注意：标注文件中的类别 ID 应从 1 开始，并按照 class_names 的顺序依次递增。如果你的标注文件不符合此要求（例如 COCO 数据集），请参考 dataset_base 中的 label_map 字段。
• 最后，在同一文件中的 yolact_edge_config 中，将 'dataset' 的值改为 'my_custom_dataset' 或你为上述配置对象所命名的名称，并将 'num_classes' 改为你的数据集中类别数加 1。之后，你就可以使用上一节中的任何训练命令了。

引用

如果你在工作中使用了本代码库，请考虑引用以下文献：

@inproceedings{yolactedge-icra2021,
  author    = {刘浩天、拉斐尔·A·里韦拉·索托、肖凡毅、李勇宰},
  title     = {YolactEdge：边缘设备上的实时实例分割},
  booktitle = {ICRA},
  year      = {2021},
}

@inproceedings{yolact-iccv2019,
  author    = {丹尼尔·博利亚、周冲、肖凡毅、李勇宰},
  title     = {YOLACT：实时实例分割},
  booktitle = {ICCV},
  year      = {2019},
}

联系方式

如有关于我们的论文或代码的问题，请联系 刘浩天 或 拉斐尔·A·里韦拉·索托。

YolactEdge 快速上手指南

YolactEdge 是首个能在小型边缘设备（如 Jetson AGX Xavier）上以实时速度运行的实例分割模型。它在保持高精度的同时，通过 TensorRT 优化实现了极高的推理帧率。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04)
• 硬件:
  • NVIDIA GPU (支持 CUDA)
  • 边缘设备：NVIDIA Jetson 系列 (AGX Xavier, Nano 等)
• 软件依赖:
  • Python 3.6+
  • PyTorch 1.7+
  • CUDA & cuDNN (版本需与 PyTorch 匹配)
  • TensorRT (用于加速，可选但强烈推荐)
  • OpenCV, Cython, pycocotools

提示: 为了简化环境配置，项目提供了官方 Dockerfile。如果您熟悉 Docker，推荐使用以下方式一键搭建完整环境：

# 构建并运行容器（参考项目根目录 docker 文件夹）
docker build -t yolact_edge .
docker run --gpus all -it yolact_edge

2. 安装步骤

步骤一：克隆代码库

git clone https://github.com/haotian-liu/yolact_edge.git
cd yolact_edge

步骤二：安装 Python 依赖

建议创建虚拟环境后安装依赖。

pip install -r requirements.txt

注：如果 pycocotools 安装失败，可尝试使用国内镜像源加速：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

步骤三：编译扩展模块

YolactEdge 包含自定义的 CUDA 内核，需要编译安装。

python setup.py build_ext --inplace

步骤四：下载预训练模型

将下载的权重文件放入 ./weights 目录。以下是基于 COCO 数据集的常用模型（含国内镜像加速链接）：

骨干网络	模型类型	下载链接 (Google Drive)	国内镜像 (OneDrive)
ResNet-101	YolactEdge (推荐)	下载	镜像
ResNet-50	YolactEdge	下载	镜像
MobileNet-V2	YolactEdge (轻量)	下载	镜像

(请将下载的文件重命名为表格中对应的文件名或直接使用命令指定路径)

3. 基本使用

安装完成后，您可以立即进行推理测试。以下命令默认使用置信度阈值 0.3 并展示前 100 个检测结果。

场景一：处理单张图片

显示结果并在窗口中弹出：

python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --image=my_image.png

保存处理后的图片到文件：

python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --image=input_image.png:output_image.png

场景二：实时视频/摄像头推理

调用摄像头（索引 0）进行实时演示：

# --video_multiframe 和 --trt_batch_size 用于提升多帧并行处理性能
python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --video_multiframe=2 --trt_batch_size 2 --video=0

处理本地视频文件：

python eval.py --trained_model=weights/yolact_edge_54_800000.pth --score_threshold=0.3 --top_k=100 --video=my_video.mp4

场景三：性能基准测试

在验证集上测试模型的 FPS 和精度：

python eval.py --trained_model=weights/yolact_edge_54_800000.pth --benchmark --max_images=1000

高级选项提示

• TensorRT 安全模式: 如果在 TensorRT 推理中遇到后处理错误，可添加 --use_tensorrt_safe_mode 参数。
• 禁用 TensorRT: 若仅需测试原始模型性能，可添加 --disable_tensorrt。
• FP16 优化: 若不使用 INT8 校准，可使用 --use_fp16_tensorrt 进行半精度加速。

查看更多参数说明：

python eval.py --help