Pytorch ReID

强大、小巧、友好

一个基于 pytorch 的小型、友好且强大的对象重识别（Object-reID）基线代码，自2017年以来。

• 强大。 它在多个顶级会议的论文中与新的基线结果一致，例如，联合判别与生成学习用于人脸重识别（CVPR19），超越部分模型：通过精炼部分池化进行人脸检索（ECCV18），用于人脸重识别的相机风格适应（CVPR18）。我们仅使用softmax loss就达到了Rank@1=88.24%，mAP=70.68%。

• 小巧。 使用fp16（由Nvidia apex支持），我们的基线仅需2GB GPU内存即可进行训练。

• 友好。 你可以使用现成的选项，只需一行代码即可应用许多最新的技巧。此外，如果你是对象重识别的新手，建议先查看我们的**教程**（阅读时间8分钟）:+1: 。

分享至

教程

• 8分钟教程，8分钟教程
• 中文视频简介
• 教程答疑

目录

• 特性
• 一些新闻
• 训练好的模型
• 先决条件
• 入门指南
  • 安装
  • 数据集准备
  • 训练
  • 测试
  • 评估
• 其他数据集训练技巧
• 如何引用
• 相关仓库

特性

目前我们支持：

训练

• 使用免费GPU在Google Colab上运行代码。点击这里（感谢@ronghao233）
• DG-Market（CVPR 2019 口头报告的Market总合成数据集，大小为市场的10倍）
• Swin Transformer / EfficientNet / HRNet
• ResNet/ResNet-ibn/DenseNet
• Circle Loss、三元损失、对比损失、球面损失、提升损失、Arcface、Cosface 和 实例损失
• 基于 apex 的Float16以节省GPU内存
• 基于部分卷积的基线（PCB）
• 随机擦除
• 线性预热
• torch.compile（更快的训练速度）
• DDP（多GPU）

测试

• TensorRT
• Pytorch JIT
• 将卷积层和批量归一化层融合为一个卷积层
• 多查询评估
• Re-Ranking（CPU和GPU版本）
• 可视化训练曲线
• 可视化排名结果
• 可视化热力图

这里我们提供生成结果所用的超参数和架构。某些参数（如学习率）远未达到最佳。请大胆更改它们并观察效果。

附：使用相似的结构，我们在 Matconvnet 上达到了 Rank@1=87.74% mAP=69.46%。（batchsize=8，dropout=0.75）你可以参考这里。不同框架需要略微不同的调优方式。

一些新闻

2024年1月12日 我们将在ACM ICMR 2024上举办多媒体对象重识别研讨会。欢迎分享您的见解。期待在泰国普吉岛见到你！😃 研讨会链接是 https://www.zdzheng.xyz/MORE2024/ 。提交截止日期是2024年4月15日。优秀论文将被推荐到 ACM TOMM 特刊（CCF-B）。 （需要重新提交。）

2023年8月12日 大型人物语言模型现已可以在这里获取，被ACM MM'23收入。欢迎查看。

2023年3月19日 我们在IEEE智能交通系统会议（ITSC）上举办了一个特别会议，涵盖对象重识别和点云话题。论文截止日期为2023年5月15日，论文通知日期为2023年6月30日。提交时请选择会话代码w7r4a。更多详情请见特别会议网站。

2023年3月9日 Market-1501变成3D了。请查看我们从单2D到3D重建的作品 https://github.com/layumi/3D-Magic-Mirror 。

2020年新闻

2020年6月11日 人们生活在3D世界中。我们发布了一个新的行人重识别代码Person Re-identification in the 3D Space，在3D空间中进行表示学习。欢迎查看。

2020年4月30日 我们将此代码应用于AICity Challenge 2020，在re-id赛道上取得了第一名提交:red_car:。查看这里。

2020年3月1日 我们发布了一个新的图像检索数据集，名为University-1652，用于无人机视角目标定位和无人机导航:helicopter:。它的设置与行人re-ID相似。欢迎查看。

python train.py --fp16

python prepare.py
python train.py
python test.py
python demo.py --query_index 777

python prepare.py
python train.py
python test.py --multi
python evaluate_gpu.py

python train.py --PCB --batchsize 64 --name PCB-64
python test.py --PCB --name PCB-64

训练模型

我重新训练了几个模型，结果可能与原始模型不同。仅供快速参考，你可以直接使用这些模型。 下载链接在这里。

• 更多的训练迭代可能会带来更好的结果。
• Swin 需要更多的 GPU 内存 (运行需要11G GPU)。
• DG-Market 的 --DG 超参数没有优化。更好的超参数可能会带来更好的结果。

不同的损失函数

我没有优化超参数。你可以自由调整它们以获得更好的性能。 |方法 | Rank@1 | mAP| 参考资料| | -------- | ----- | ---- | ---- | | CE | 92.01% | 79.31% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_100 --total 100 ; python test.py --name warm5_s1_b32_lr8_p0.5_100| | CE + Sphere [论文] | 92.01% | 79.39% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_sphere100 --sphere --total 100; python test.py --name warm5_s1_b32_lr8_p0.5_sphere100 | | CE + Triplet [论文] | 92.40% | 79.71% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_triplet100 --triplet --total 100; python test.py --name warm5_s1_b32_lr8_p0.5_triplet100 | | CE + Lifted [论文]| 91.78% | 79.77% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_lifted100 --lifted --total 100; python test.py --name warm5_s1_b32_lr8_p0.5_lifted100 | | CE + Instance [论文] | 92.73% | 81.11% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_instance100_gamma64 --instance --ins_gamma 64 --total 100 ; python test.py --name warm5_s1_b32_lr8_p0.5_instance100_gamma64| | CE + Contrast [论文] | 92.28% | 81.42% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_contrast100 --contrast --total 100; python test.py --name warm5_s1_b32_lr8_p0.5_contrast100| | CE + Circle [论文] | 92.46% | 81.70% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_circle100 --circle --total 100 ; python test.py --name warm5_s1_b32_lr8_p0.5_circle100 | | CE + Contrast + Sphere | 92.79% | 82.02% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 32 --lr 0.08 --name warm5_s1_b32_lr8_p0.5_cs100 --contrast --sphere --total 100; python test.py --name warm5_s1_b32_lr8_p0.5_cs100| | CE + Contrast + Triplet (Long) | 92.61% | 82.01% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 24 --lr 0.062 --name warm5_s1_b24_lr6.2_p0.5_contrast_triplet_133 --contrast --triplet --total 133 ; python test.py --name warm5_s1_b24_lr6.2_p0.5_contrast_triplet_133 | | CE + Contrast + Circle (Long) | 92.19% | 82.07% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 24 --lr 0.08 --name warm5_s1_b24_lr8_p0.5_contrast_circle133 --contrast --circle --total 133 ; python test.py --name warm5_s1_b24_lr8_p0.5_contrast_circle133 | | CE + Contrast + Sphere (Long) | 92.84% | 82.37% | python train.py --warm_epoch 5 --stride 1 --erasing_p 0.5 --batchsize 24 --lr 0.06 --name warm5_s1_b24_lr6_p0.5_contrast_sphere133 --contrast --sphere --total 133 ; python test.py --name warm5_s1_b24_lr6_p0.5_contrast_sphere133 |

模型结构

你可以从 model.py 中了解更多信息。 我们添加了一个线性层（瓶颈层）、一个batchnorm层和relu。

先决条件

• Python 3.6+
• GPU 显存 >= 6G
• Numpy
• Pytorch 0.3+
• 使用 Pytorch >1.7.0 版本的 Swin-Transformer，需要安装 timm pip install timm
• 通过 pip install pretrainedmodels 安装 pretrainedmodels
• [可选] apex (用于 float16)
• [可选] pretrainedmodels

（一些报告发现更新 numpy 可以达到正确的准确率。如果你只有 50~80 的 Top1 准确率，可以尝试一下。） 我们已经成功运行了基于 numpy 1.12.1 和 1.13.1 的代码。

入门指南

安装

• 从 http://pytorch.org/ 安装 Pytorch
• 从源代码安装 Torchvision

git clone https://github.com/pytorch/vision
cd vision
python setup.py install

• [可选] 你可以跳过这个步骤。从源代码安装 apex

git clone https://github.com/NVIDIA/apex.git
cd apex
python setup.py install --cuda_ext --cpp_ext

因为 pytorch 和 torchvision 是正在进行的项目。

注意：我们的代码测试基于 Pytorch 0.3.0/0.4.0/0.5.0/1.0.0 和 Torchvision 0.2.0/0.2.1。

数据集 & 准备工作

下载 Market1501 数据集 [Google] [Baidu] 或者使用命令行:

pip install gdown 
pip install --upgrade gdown #!!重要!!
gdown 0B8-rUzbwVRk0c054eEozWG9COHM

准备工作：将相同id的图片放在一个文件夹中。你可以使用

python prepare.py

记得将数据集路径更改为你自己的路径。

此外，你也可以在 [DukeMTMC-reID 数据集]( GoogleDriver 或 (BaiduYun 密码: bhbh)) 上测试我们的代码。或者使用命令行：

gdown 1jjE85dRCMOgRtvJ5RQV9-Afs-2_5dY3O

我们的基线代码在 DukeMTMC-reID 上的性能不是很高 Rank@1=64.23%, mAP=43.92%。超参数需要调优。

• [可选] DG-Market 是一个包含 128,307 张图片的生成行人数据集，用于训练一个鲁棒模型。

训练

通过以下命令训练模型

python train.py --gpu_ids 0 --name ft_ResNet50 --train_all --batchsize 32  --data_dir your_data_path

--gpu_ids 指定运行的 GPU。

--name 模型名。

--data_dir 训练数据的路径。

--train_all 使用所有图片进行训练。

--batchsize 批次大小。

--erasing_p 随机擦除概率。

通过以下命令进行随机擦除训练

python train.py --gpu_ids 0 --name ft_ResNet50 --train_all --batchsize 32  --data_dir your_data_path --erasing_p 0.5

如果你想使用多GPU，建议使用 DDP (train_DDP.py) 而不是 DP (train.py)。因为 DP 缺少 torch 的支持，可能会遇到一些 NaN 的情况。 你可以通过运行 DDP.sh 来调用 train_DDP.py。

bash DDP.sh 

测试

通过以下命令使用训练好的模型提取特征

python test.py --gpu_ids 0 --name ft_ResNet50 --test_dir your_data_path  --batchsize 32 --which_epoch 59

--gpu_ids 指定运行的 GPU。

--batchsize 批次大小。

--name 训练模型的目录名。 --which_epoch 选择第 i 个模型。

--data_dir 测试数据的路径。

评估

python evaluate.py

它会输出 Rank@1、Rank@5、Rank@10 和 mAP 结果。您也可以尝试 evaluate_gpu.py 以使用 GPU 进行更快的评估。

关于 mAP 计算，您可以参考 Oxford Building 的 C++ 代码。我们使用三角 mAP 计算方式（与 Market1501 的原始代码一致）。

重排名

python evaluate_rerank.py

运行可能需要超过 10G 内存。 因此，如果可能，请在性能强大的机器上运行。

它会输出 Rank@1、Rank@5、Rank@10 和 mAP 结果。

提示

注意相机 ID 的格式和相机数量。

对于某些数据集，例如 MSMT17，有超过 10 个相机。您需要修改 prepare.py 和 test.py 来读取两位数的相机 ID。

对于一些车辆重识别数据集，例如 VeRi，您也需要修改 prepare.py 和 test.py。它有不同的命名规则。 https://github.com/layumi/Person_reID_baseline_pytorch/issues/107 （抱歉，这页面是中文的）

引用

以下论文使用并报告了基线模型的结果。您可以在您的论文中引用它。

@article{zheng2019joint,
  title={Joint discriminative and generative learning for person re-identification},
  author={Zheng, Zhedong and Yang, Xiaodong and Yu, Zhiding and Zheng, Liang and Yang, Yi and Kautz, Jan},
  journal={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  year={2019}
}

以下论文可能是最早使用瓶颈基线的两篇论文。您可以在您的论文中引用它们。

@article{DBLP:journals/corr/SunZDW17,
  author    = {Yifan Sun and
               Liang Zheng and
               Weijian Deng and
               Shengjin Wang},
  title     = {SVDNet for Pedestrian Retrieval},
  booktitle   = {ICCV},
  year      = {2017},
}

@article{hermans2017defense,
  title={In Defense of the Triplet Loss for Person Re-Identification},
  author={Hermans, Alexander and Beyer, Lucas and Leibe, Bastian},
  journal={arXiv preprint arXiv:1703.07737},
  year={2017}
}

基础模型

@article{zheng2018discriminatively,
  title={A discriminatively learned CNN embedding for person reidentification},
  author={Zheng, Zhedong and Zheng, Liang and Yang, Yi},
  journal={ACM Transactions on Multimedia Computing, Communications, and Applications (TOMM)},
  volume={14},
  number={1},
  pages={13},
  year={2018},
  publisher={ACM}
}

@article{zheng2020vehiclenet,
  title={VehicleNet: Learning Robust Visual Representation for Vehicle Re-identification},
  author={Zheng, Zhedong and Ruan, Tao and Wei, Yunchao and Yang, Yi and Mei, Tao},
  journal={IEEE Transaction on Multimedia (TMM)},
  year={2020}
}

相关仓库

1. Pedestrian Alignment Network
2. 2stream Person re-ID
3. Pedestrian GAN
4. Language Person Search
5. DG-Net
6. 3D Person re-ID