MeMOTR

MeMOTR: 基于长期记忆增强transformer的多目标跟踪（ICCV 2023）的官方实现。

作者：高若鹏，王利民。

MeMOTR是一个完全端到端的基于Transformer的记忆增强多目标跟踪器。我们利用长期记忆注入和定制的记忆注意力层，显著提高了关联性能。

新闻 :fire:

• 2024.05.09：我们发布了MOTIP，一种将多目标跟踪任务视为ID预测问题的新视角 :telescope:。

• 2024.02.21：我们在arxiv版本（补充部分）中添加了SportsMOT的结果。如果您在SportsMOT比较中引用我们的跟踪器，我们将不胜感激 :chart_with_upwards_trend:。

• 2023.12.24：我们发布了BDD100K上的代码、脚本和检查点 :car:。

• 2023.12.13：我们实现了一个jupyter notebook，可以在您自己的视频上运行我们的模型 :movie_camera:。

• 2023.11.07：我们发布了SportsMOT上的脚本和检查点 :basketball:。

• 2023.08.24：我们发布了DanceTrack上的脚本和检查点 :dancer:。

• 2023.08.09：我们发布了主要代码。更多配置、脚本和检查点将很快发布 :soon:。

安装

conda create -n MeMOTR python=3.10  # 创建虚拟环境
# 我记得我使用了Python 3.10的一些新特性，但我不确定具体是哪些。
conda activate MeMOTR               # 激活环境
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia
# 我们的代码主要在PyTorch 1.13.1上运行，
# 但它应该也兼容早期的PyTorch版本（如1.12.1）。
# 然而，太早的pytorch版本可能会导致一些需要修复的问题，因为我们使用了pytorch的一些新提出的特性（例如，ResNet50_Weights）。
conda install matplotlib pyyaml scipy tqdm tensorboard
pip install opencv-python

您还需要编译Deformable Attention CUDA算子：

# 来自 https://github.com/fundamentalvision/Deformable-DETR
cd ./models/ops/
sh make.sh
# 如果需要，您可以测试这些算子：
python test.py

数据

您应该将解压后的MOT17和CrowdHuman数据集分别放入DATADIR/MOT17/images/和DATADIR/CrowdHuman/images/。然后通过运行相应的脚本生成真实标注文件：./data/gen_mot17_gts.py和./data/gen_crowdhuman_gts.py。

最后，您应该得到以下数据集结构：

DATADIR/
  ├── DanceTrack/
  │ ├── train/
  │ ├── val/
  │ ├── test/
  │ ├── train_seqmap.txt
  │ ├── val_seqmap.txt
  │ └── test_seqmap.txt
  ├── SportsMOT/
  │ ├── train/
  │ ├── val/
  │ ├── test/
  │ ├── train_seqmap.txt
  │ ├── val_seqmap.txt
  │ └── test_seqmap.txt
  ├── MOT17/
  │ ├── images/
  │ │ ├── train/     # 从MOT17解压
  │ │ └── test/      # 从MOT17解压
  │ └── gts/
  │   └── train/     # 由./data/gen_mot17_gts.py生成
  └── CrowdHuman/
    ├── images/
    │ ├── train/     # 从CrowdHuman解压
    │ └── val/       # 从CrowdHuman解压
    └── gts/
      ├── train/     # 由./data/gen_crowdhuman_gts.py生成
      └── val/       # 由./data/gen_crowdhuman_gts.py生成

预训练

我们使用官方的DAB-Deformable-DETR（使用R50骨干网络）在COCO数据集上预训练的权重初始化我们的模型，您也可以在这里下载我们使用的检查点。然后将检查点放在项目根目录下。

DanceTrack脚本

训练

使用8个GPU在DanceTrack上训练MeMOTR（建议使用>=32 GB内存的GPU，如V100-32GB或其他）：

python -m torch.distributed.run --nproc_per_node=8 main.py --use-distributed --config-path ./configs/train_dancetrack.yaml --outputs-dir ./outputs/memotr_dancetrack/ --batch-size 1 --data-root <您的数据目录路径>

如果您的GPU内存小于32 GB，我们还实现了一个内存优化版本（通过运行选项--use-checkpoint），如论文中讨论的那样，我们使用梯度检查点来减少分配的GPU内存。以下训练脚本只需要约10 GB的GPU内存：

python -m torch.distributed.run --nproc_per_node=8 main.py --use-distributed --config-path ./configs/train_dancetrack.yaml --outputs-dir ./outputs/memotr_dancetrack/ --batch-size 1 --data-root <您的数据目录路径> --use-checkpoint

提交和评估

您可以使用此脚本在DanceTrack验证集上评估训练好的模型：

python main.py --mode eval --data-root <您的数据目录路径> --eval-mode specific --eval-model <检查点文件名> --eval-dir ./outputs/memotr_dancetrack/ --eval-threads <您的GPU数量>

对于提交，您可以使用以下脚本：

python -m torch.distributed.run --nproc_per_node=8 main.py --mode submit --submit-dir ./outputs/memotr_dancetrack/ --submit-model <检查点文件名> --use-distributed --data-root <您的数据目录路径>

此外，如果您只想直接使用我们训练好的检查点进行评估或提交，您可以在这里获取我们在论文中使用的检查点。然后将此检查点放入./outputs/memotr_dancetrack/并运行上述脚本。

MOT17脚本

提交

对于提交，您可以使用以下脚本：

python -m torch.distributed.run --nproc_per_node=8 main.py --mode submit --config-path ./outputs/memotr_mot17/train/config.yaml --submit-dir ./outputs/memotr_mot17/ --submit-model <检查点文件名> --use-distributed --data-root <你的数据目录路径>

另外，你可以直接在这里下载我们训练好的检查点。然后将其放入 ./outputs/memotr_mot17/ 并运行上述脚本进行提交，以获取MOT17测试集的提交文件。

SportsMOT和其他数据集的脚本

你可以替换DanceTrack脚本中的--config-path。例如，将./configs/train_dancetrack.yaml改为./configs/train_sportsmot.yaml以在SportsMOT上进行训练。

结果

DanceTrack测试集上的多目标跟踪

方法	HOTA	DetA	AssA	检查点
MeMOTR	68.5	80.5	58.4	Google Drive
MeMOTR (Deformable DETR)	63.4	77.0	52.3	Google Drive

SportsMOT测试集上的多目标跟踪

对于所有实验，我们不使用额外数据（如CrowdHuman）进行训练。

方法	HOTA	DetA	AssA	检查点
MeMOTR	70.0	83.1	59.1	Google Drive
MeMOTR (Deformable DETR)	68.8	82.0	57.8	Google Drive

MOT17测试集上的多目标跟踪

方法	HOTA	DetA	AssA	检查点
MeMOTR	58.8	59.6	58.4	Google Drive

BDD100K验证集上的多类别多目标跟踪

方法	mTETA	mLocA	mAssocA	检查点
MeMOTR	53.6	38.1	56.7	Google Drive

联系方式

• 高若鹏: ruopenggao@gmail.com

引用

@InProceedings{MeMOTR,
    author    = {Gao, Ruopeng and Wang, Limin},
    title     = {{MeMOTR}: Long-Term Memory-Augmented Transformer for Multi-Object Tracking},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2023},
    pages     = {9901-9910}
}

星标

致谢

• Deformable DETR
• DAB DETR
• MOTR
• TrackEval
• CV-Framework
• Star History