环境

  • 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,默认可训练的车辆数据集包括VeRiVehicleIDVERIWild,应用场景非常广泛。

架构

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

大模型开源的embedding_机器学习


fast-reid

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

训练部分包括

  1. Pre-processing即预处理,包括各种数据增强方法,如ResizeFlippingRandom erasingAuto-augmentRandom patchCutout
  2. Backbone即骨干网,如ResNetResNestResNeXtIBN(Instance Batch Normalization)Non-local
  3. Aggregation即聚合,将骨干网生成的特征聚合成一个全局特征,如AttentionGem PoolingAvg PoolingMax pooling
  4. Head模块,用于对生成的全局特征进行归一化、纬度约减
  5. Loss部分,包括Cross-entropy lossTriplet lossArcface lossCircle loss
  6. 另外在训练策略方面,使用了包含Learning ratewarm-upBackbone freezeConsine decay

推理部分包括

  1. Distance Metric距离度量,支持欧式距离、余弦和Deep Spatial Reconstruction(DSR)
  2. Post processing后处理,K-reciprocalQuery 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-1501DukeMTMCMSMT17VehicleIDVeRiVERIWild。数据集默认存放在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

大模型开源的embedding_人工智能_02


fast-reid

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

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

大模型开源的embedding_人工智能_03


fast-reid

训练结果可视化

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,希望可以合并进去。然后,我们再次执行上述脚本

大模型开源的embedding_机器学习_04


fast-reid

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

大模型开源的embedding_机器学习_05


fast-reid

大模型开源的embedding_tensorflow_06


fast-reid

转换成其它格式

FastReID提供了一系列的python脚本,帮你快速实现从fastreid的模型转换成其它格式的模型,如CaffeONNXTRT

大模型开源的embedding_机器学习_07


fast-reid

这里我们以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

大模型开源的embedding_深度学习_08


fast-reid

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

大模型开源的embedding_深度学习_09


fast-reid

接下来执行

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

大模型开源的embedding_大模型开源的embedding_10


fast-reid

参考资料