MEGA 用于视频目标检测

作者：Yihong Chen、Yue Cao、Han Hu、Liwei Wang。

本仓库是 CVPR 2020 接收的论文 "Memory Enhanced Global-Local Aggregation for Video Object Detection" 的官方实现。该仓库基于 maskrcnn_benchmark 提供了我们方法 MEGA 的 PyTorch 实现，并包含一些训练脚本，用于复现我们在论文中报告的 ImageNet VID 数据集上的实验结果。

此外，本仓库还实现了其他几种算法，如 FGFA 和 RDN。欢迎添加新的方法！期待您的 Pull Request！我们希望这个仓库能够帮助推动视频目标检测及相关领域的进一步研究。 :)

引用 MEGA

如果我们的工作对您的研究有所帮助，请在您的出版物中引用我们的论文：

@inproceedings{chen20mega,
    Author = {Chen, Yihong and Cao, Yue and Hu, Han and Wang, Liwei},
    Title = {Memory Enhanced Global-Local Aggregation for Video Object Detection},
    Conference = {CVPR},
    Year = {2020}
}

更新

• 添加了针对 motion-IoU 的 AP 评估代码。仅适用于 ImageNet VID 数据集。（2020年6月19日）
• 增加了可视化演示（支持图片文件夹和视频）。（2020年6月17日）
• 添加了使用 ResNet-50 主干网络的结果。（2020年4月13日）
• 发布了 Deep Feature Flow 的代码和预训练权重。（2020年3月30日）

主要结果

预训练模型现已在百度网盘（提取码：neck）和 Google Drive 上提供。

模型	主干网络	AP50	AP (fast)	AP (med)	AP (slow)	链接
单帧基线	ResNet-101	76.7	52.3	74.1	84.9	Google
DFF	ResNet-101	75.0	48.3	73.5	84.5	Google
FGFA	ResNet-101	78.0	55.3	76.9	85.6	Google
RDN-base	ResNet-101	81.1	60.2	79.4	87.7	Google
RDN	ResNet-101	81.7	59.5	80.0	89.0	Google
MEGA	ResNet-101	82.9	62.7	81.6	89.4	Google

模型	主干网络	AP50	AP (fast)	AP (med)	AP (slow)	链接
单帧基线	ResNet-50	71.8	47.2	69.2	80.6	Google
DFF	ResNet-50	70.4	43.6	68.9	80.8	Google
FGFA	ResNet-50	74.3	50.6	72.3	84.0	Google
RDN-base	ResNet-50	76.7	53.8	74.8	85.4	Google
MEGA	ResNet-50	77.3	56.5	75.7	85.2	Google

注意：ResNet-50 主干网络的性能不太稳定。

注意：针对 motion-IoU 的 AP 评估代码与 FGFA 中的原始实现略有不同。我认为原始实现确实有些奇怪，因此我对其进行了修改。所以这些结果可能无法直接与 FGFA 及其他使用 MXNet 版本评估代码的方法的结果进行比较。不过，在相同的评估协议下，我们仍然可以判断哪种方法相对更好。

安装

请按照 INSTALL.md 中的说明进行安装。

数据准备

请从 这里 下载 ILSVRC2015 DET 和 ILSVRC2015 VID 数据集。之后，建议将数据集路径符号链接到 datasets/ 目录下。路径结构应如下所示：

./datasets/ILSVRC2015/
./datasets/ILSVRC2015/Annotations/DET
./datasets/ILSVRC2015/Annotations/VID
./datasets/ILSVRC2015/Data/DET
./datasets/ILSVRC2015/Data/VID
./datasets/ILSVRC2015/ImageSets

注意：我们已经在 datasets/ILSVRC2015/ImageSets 目录下提供了用于训练和测试模型的所有图像列表，格式为 txt 文件。您无需更改它们。

使用

注意：首次运行本项目时会生成缓存文件，这可能需要一些时间！请不必担心！

注意：目前，一张 GPU 只能处理一张图像。请勿在同一张 GPU 上放置两张或更多图像！

注意 我们提供了名为 BASE_RCNN_{}gpus.yaml 的模板文件，它会自动调整批次大小和其他相关设置。这种行为类似于 detectron2。如果您想使用不同数量的 GPU 进行训练，请自行修改 :) 但请务必确保 每张 GPU 只处理一张图像！也就是说，您应始终将 SOLVER.IMS_PER_BATCH 和 TEST.IMS_PER_BATCH 设置为所使用的 GPU 数量。

推理

在验证数据集上进行测试的推理命令行：

python -m torch.distributed.launch \
    --nproc_per_node 4 \
    tools/test_net.py \
    --config-file configs/MEGA/vid_R_101_C4_MEGA_1x.yaml \
    --motion-specific \
    MODEL.WEIGHT MEGA_R_101.pth 
    

请注意：

1. 如果您的模型名称不同，请将 MEGA_R_101.pth 替换为您自己的模型文件名。
2. 如果您想评估不同的模型，请将 --config-file 更改为相应配置文件，并将 MODEL.WEIGHT 更改为该模型的权重文件。
3. 如果您不想评估运动 IoU 特定的 AP，只需删除 --motion-specific 即可。
4. 测试过程较为耗时，请耐心等待！
5. 由于在超过 17 万帧的数据上进行测试非常耗时，我们提供了直接基于生成的边界框进行测试的功能，这些结果会自动保存到训练目录下的 predictions.pth 文件中。这意味着您无需每次都从头开始运行评估。您可以通过以下命令访问该文件并进行评估：

    python tools/test_prediction.py \
        --config-file configs/MEGA/vid_R_101_C4_MEGA_1x.yaml \
        --prediction [您使用 MEGA 生成的 predictions.pth 文件]
        --motion-specific

训练

以下命令行将在 4 张 GPU 上使用同步随机梯度下降（SGD）训练 MEGA_R_101_FPN_1x 模型：

python -m torch.distributed.launch \
    --nproc_per_node=4 \
    tools/train_net.py \
    --master_port=$((RANDOM + 10000)) \
    --config-file configs/MEGA/vid_R_101_C4_MEGA_1x.yaml \
    --motion-specific \
    OUTPUT_DIR training_dir/MEGA_R_101_1x
    

请注意：

1. 模型将被保存到 OUTPUT_DIR 目录中。
2. 如果您希望使用其他骨干网络训练 MEGA 或其他方法，请更改 --config-file。
3. 对于 FGFA 和 DFF 的训练，我们需要 FlowNet 的预训练权重。我们提供了一个转换后的版本在此。下载后，请将其放置在 models/ 目录下。更多细节请参阅 config/defaults.py 文件及代码。
4. 对于 RDN 的训练，我们采用与其原始论文中描述相同的两阶段训练策略。第一阶段应使用配置文件 configs/RDN/vid_R_101_C4_RDN_base_1x.yaml 进行训练。第二阶段则需将 MODEL.WEIGHT 设置为第一阶段训练结束时生成的最终模型文件名。或者您可以将该模型文件重命名为 RDN_base_R_101.pth 并将其放入 models/ 目录中，然后直接使用配置文件 configs/RDN/vid_R_101_C4_RDN_1x.yaml 进行第二阶段训练。
5. 如果您不希望在训练结束时评估运动 IoU 特定的 AP，只需删除 --motion-specific 即可。

示例用法

请参考 demo/README.md，了解如何可视化您自己的图片或视频。

自定义

如果您希望在自己的数据集上使用这些方法，或实现您自己的新方法，请参考 CUSTOMIZE.md。

参与项目贡献

欢迎任何拉取请求或问题反馈。

MEGA.pytorch 快速上手指南

MEGA (Memory Enhanced Global-Local Aggregation) 是一个基于 PyTorch 的视频目标检测开源项目，由 CVPR 2020 论文提出。本指南帮助开发者快速搭建环境并运行模型。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 16.04/18.04)
• Python: 3.6 或更高版本
• PyTorch: 1.4 或更高版本 (需与 CUDA 版本匹配)
• GPU: 支持 CUDA 的 NVIDIA 显卡 (训练时建议多卡)
• 其他依赖: torchvision, ninja, yacs, cython, opencv-python 等

注意：本项目基于 maskrcnn_benchmark 架构开发，安装过程较为特殊，请严格按照下方步骤操作。

安装步骤

1. 克隆代码库

git clone https://github.com/scalsol/mega.pytorch.git
cd mega.pytorch

2. 创建虚拟环境并安装基础依赖

建议使用 Conda 管理环境：

conda create -n mega python=3.7
conda activate mega
conda install pytorch==1.4.0 torchvision==0.5.0 cudatoolkit=10.1 -c pytorch
pip install ninja yacs cython matplotlib opencv-python tqdm

3. 安装 maskrcnn_benchmark 核心库

进入项目目录并执行安装脚本（这将编译 CUDA 扩展）：

export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0" # 根据您的显卡算力调整，如 RTX 2080Ti 设为 7.0
python setup.py build develop

注：如果编译报错，请检查 TORCH_CUDA_ARCH_LIST 是否包含您显卡对应的算力值。

4. 数据集准备

下载 ILSVRC2015 DET 和 ILSVRC2015 VID 数据集。下载完成后，建立软链接至项目根目录下的 datasets/ 文件夹，结构如下：

./datasets/ILSVRC2015/
├── Annotations/DET
├── Annotations/VID
├── Data/DET
├── Data/VID
└── ImageSets

项目已内置所需的训练/测试列表文件 (ImageSets 目录下)，无需手动修改。

5. 下载预训练模型

从 Google Drive 或百度网盘下载预训练权重，并放置在项目根目录或指定路径。

• MEGA (ResNet-101): Google Drive | 百度网盘 (提取码: neck)

基本使用

推理测试 (Inference)

使用下载的预训练模型在验证集上进行测试。以下命令示例使用 4 张 GPU：

python -m torch.distributed.launch \
    --nproc_per_node 4 \
    tools/test_net.py \
    --config-file configs/MEGA/vid_R_101_C4_MEGA_1x.yaml \
    --motion-specific \
    MODEL.WEIGHT MEGA_R_101.pth

关键参数说明：

• --nproc_per_node: 使用的 GPU 数量。重要约束：目前每个 GPU 仅能处理 1 张图像，因此 Batch Size 必须等于 GPU 数量。
• --motion-specific: 启用针对运动 IoU 的特定评估指标（可选，若不需要可删除此参数）。
• MODEL.WEIGHT: 指向您的 .pth 权重文件路径。

提示：全量测试非常耗时。首次运行后生成的边界框结果会自动保存为 predictions.pth。后续若只需重新评估指标而无需重新推理，可使用以下命令加速：

python tools/test_prediction.py \
    --config-file configs/MEGA/vid_R_101_C4_MEGA_1x.yaml \
    --prediction [YOUR predictions.pth] \
    --motion-specific

模型训练 (Training)

使用 4 张 GPU 训练 MEGA 模型 (ResNet-101 backbone)：

python -m torch.distributed.launch \
    --nproc_per_node=4 \
    tools/train_net.py \
    --master_port=$((RANDOM + 10000)) \
    --config-file configs/MEGA/vid_R_101_C4_MEGA_1x.yaml \
    --motion-specific \
    OUTPUT_DIR training_dir/MEGA_R_101_1x

注意事项：

1. 模型检查点将保存在 OUTPUT_DIR 指定的目录中。
2. 若训练 FGFA 或 DFF 方法，需预先下载 FlowNet 的预训练权重并放入 models/ 目录。
3. 若训练 RDN 方法，需采用两阶段训练策略（先训练 base 版本，再加载权重进行第二阶段训练）。
4. 配置文件 BASE_RCNN_{}gpus.yaml 会自动调整 Batch Size，但请务必保证 SOLVER.IMS_PER_BATCH 与实际使用的 GPU 数量一致（即 1 GPU = 1 Image）。

可视化演示 (Demo)

若需对自己的图片或视频进行可视化演示，请参考项目内 demo/README.md 文档进行操作。