强大的国产ReID开源工具

环境

• ubuntu 18.04 64位
• GTX 1070Ti
• anaconda with python 3.7
• pytorch 1.6.0 + CUDA 10.1
• fast-reid 1.3.0

前言

FastReID 是京东 AI 研究院推出的一个 SOTA 级的 ReID 算法研究平台，也是迄今为止最完整的、性能最高的 ReID 工具箱，包括模型训练、模型评估、模型微调和模型部署。训练部分不仅支持单卡和多卡 GPU 训练，还支持多机训练(multiple machines)。

FastReID 不但可以进行行人的 ReID，同时也可以进行车辆等其它物体的 ReID，默认可训练的车辆数据集包括 VeRi、VehicleID 和 VERIWild，应用场景非常广泛。

架构

这是一张从其论文中截下来的系统架构图

上图中包含2部分内容，训练部分和推理部分

训练部分包括

1. Pre-processing即预处理，包括各种数据增强方法，如Resize、Flipping、Random erasing、Auto-augment、Random patch、Cutout等

2. Backbone即骨干网，如ResNet、ResNest、ResNeXt、IBN(Instance Batch Normalization)、Non-local

3. Aggregation即聚合，将骨干网生成的特征聚合成一个全局特征，如Attention、Gem Pooling、Avg Pooling、Max pooling

4. Head模块，用于对生成的全局特征进行归一化、纬度约减

5. Loss部分，包括Cross-entropy loss、Triplet loss、Arcface loss、Circle loss

6. 另外在训练策略方面，使用了包含Learning rate、warm-up、Backbone freeze、Consine decay

推理部分包括

1. Distance Metric 距离度量，支持欧式距离、余弦和 Deep Spatial Reconstruction(DSR)

2. Post processing 后处理，K-reciprocal 和 Query Expansion(QE)

复现

创建个全新的虚拟环境

conda create -n pytorch1.6 python=3.7
conda activate pytorch1.6

接着去拉取源码，目前的稳定版是v1.3.0

wget https://github.com/JDAI-CV/fast-reid/archive/refs/tags/v1.3.0.zip
unzip v1.3.0.zip
cd fast-reid-1.3.0

最后安装下其它的依赖

# 如果没有gpu的话，就pip install torch torchvision即可
pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
pip install -r docs/requirements.txt

# 使用cython加速
pip install cython
cd fastreid/evaluation/rank_cylib
make all

训练及验证

FastReID 已经原生支持多种数据集，像 Market-1501、DukeMTMC、MSMT17、VehicleID、VeRi 和 VERIWild。数据集默认存放在 dataset 文件夹，当然也可以通过环境变量 FASTREID_DATASETS 来指定其它目录，或者通过软链接的方式链接到真正的数据目录。

这里我们使用 Market-1501 数据集为例训练一下，首先下载数据集

链接：https://pan.baidu.com/s/1i9aiZx-EC3fjhn3uWTKZjw
提取码：up8x

下载后解压到目录 dataset 下，目录结构是这样的

(pytorch1.6) xugaoxiang@1070Ti:~/workshop/fast-reid-1.3.0/datasets$ tree -d
.
└── Market-1501-v15.09.15
    ├── bounding_box_test
    ├── bounding_box_train
    ├── gt_bbox
    ├── gt_query
    └── query

6 directories

然后就可以训练了，使用命令

python tools/train_net.py --config-file ./configs/Market1501/bagtricks_R50.yml

针对训练出来的模型进行评估

python tools/train_net.py --config-file ./configs/Market1501/bagtricks_R50.yml --eval-only MODEL.WEIGHTS "logs/market1501/bagtricks_R50/model_final.pth"

训练结果可视化

python ./demo/visualize_result.py --config-file "configs/Market1501/AGW_R50.yml" --vis-label --dataset-name 'Market1501' --output 'logs/market1501/agw_R50/agw_market1501_vis' --opts MODEL.WEIGHTS "logs/market1501/agw_R50/model_final.pth"

这里报了个错

Traceback (most recent call last):
  File "demo/visualize_result.py", line 127, in <module>
    distmat = distmat.numpy()
TypeError: can't convert cuda:0 device type tensor to numpy. Use Tensor.cpu() to copy the tensor to host memory first.

我们修改 demo/visualize_result.py 中的L127，将

distmat = distmat.numpy()

改为

distmat = distmat.cpu().numpy()

我已经向官方仓库提交了 pr，希望可以合并进去。然后，我们再次执行上述脚本

结果保存在 logs/market1501/agw_R50/agw_market1501_vis 下

转换成其它格式

FastReID 提供了一系列的 python 脚本，帮你快速实现从 fastreid 的模型转换成其它格式的模型，如 Caffe、ONNX、TRT等

这里我们以 onnx 的转换为例，看看具体的步骤

python tools/deploy/onnx_export.py --config-file configs/Market1501/bagtricks_R50.yml --name baseline_R50 --output onnx_model --opts MODEL.WEIGHTS logs/market1501/bagtricks_R50/model_final.pth

转换好了之后，我们准备同一个人的若干张图片，放在input文件夹下

接下来执行

python tools/deploy/onnx_inference.py --model-path onnx_model/baseline_R50.onnx --input input/*.jpg --output onnx_output

参考资料

• https://github.com/JDAI-CV/fast-reid
• https://arxiv.org/pdf/2006.02631.pdf
• https://xugaoxiang.com/2021/03/10/market1501-yolov5-deepsort/