Reid strong baseline知识蒸馏【附代码】

本项目是在Reid strong baseline基础上进行的更新，实现的知识蒸馏。项目暂未加入目标检测部分，后期会不定时更新，请持续关注。

本相比Reid所用数据集为Markt1501，支持Resnet系列作为训练的baseline网络。训练采用表征学习+度量学习的方式，蒸馏特征蒸馏【暂未更新逻辑蒸馏】。

项目说明

1.仅支持markt1501数据集

2.支持resnet网络系列训练(未蒸馏)

3.支持resnet系列特征自蒸馏训练(暂为更新离线蒸馏和逻辑蒸馏)

4.可自由选择教师网络和学生网络(本项目的特点)

5.支持评价指标的测试

---

目录

Reid训练

知识蒸馏训练

教师网络&学生网络选择

测试

Reid相关资料学习链接

后期计划更新

---

度量学习采用三元组损失函数

数据集：mark1501(将数据集mark1501放在data文件夹下)

baseline网络：支持Resnet系列，例如resnet18、resnet34、rensnet_ibn等

Reid训练

python tools/train.py --model_name resnet50_ibn_a --model_path weights/ReID_resnet50_ibn_a.pth --IMS_PER_BATCH 8 --TEST_IMS_PER_BATCH 4 --MAX_EPOCHS 120

model_name:可支持的baseline网络

支持：resnet18,resnet34,resnet50,resnet101,resnet50_ibn_a

接着会出现下面的内容：

=> Market1501 loaded
Dataset statistics:
  ----------------------------------------
  subset   | # ids | # images | # cameras
  ----------------------------------------
  train    |   751 |    12936 |         6
  query    |   750 |     3368 |         6
  gallery  |   751 |    15913 |         6
  ----------------------------------------
  
2023-05-15 14:30:55.603 | INFO     | engine.trainer:log_training_loss:119 - Epoch[1] Iteration[227/1484] Loss: 6.767, Acc: 0.000, Base Lr: 3.82e-05
2023-05-15 14:30:55.774 | INFO     | engine.trainer:log_training_loss:119 - Epoch[1] Iteration[228/1484] Loss: 6.761, Acc: 0.000, Base Lr: 3.82e-05
2023-05-15 14:30:55.946 | INFO     | engine.trainer:log_training_loss:119 - Epoch[1] Iteration[229/1484] Loss: 6.757, Acc: 0.000, Base Lr: 3.82e-05
2023-05-15 14:30:56.134 | INFO     | engine.trainer:log_training_loss:119 - Epoch[1] Iteration[230/1484] Loss: 6.760, Acc: 0.000, Base Lr: 3.82e-05
2023-05-15 14:30:56.305 | INFO     | engine.trainer:log_training_loss:119 - Epoch[1] Iteration[231/1484] Loss: 6.764, Acc: 0.000, Base Lr: 3.82e-05
​

每个epoch训练完成后会测试一次mAP：

我这里第一个epoch的mAP达到75.1%，Rank-1:91.7%, Rank-5:97.2%, Rank-10:98.2%。

测试完成后会在log文件下保存一个pth权重，名称为mAPxx.pth，也是用该权重进行测试。

2023-05-15 14:35:59.753 | INFO     | engine.trainer:print_times:128 - Epoch 1 done. Time per batch: 261.820[s] Speed: 45.4[samples/s]
2023-05-15 14:35:59.755 | INFO     | engine.trainer:print_times:129 - ----------
The test feature is normalized
2023-05-15 14:39:51.025 | INFO     | engine.trainer:log_validation_results:137 - Validation Results - Epoch: 1
2023-05-15 14:39:51.048 | INFO     | engine.trainer:log_validation_results:140 - mAP:75.1%
2023-05-15 14:39:51.051 | INFO     | engine.trainer:log_validation_results:142 - CMC curve, Rank-1  :91.7%
2023-05-15 14:39:51.051 | INFO     | engine.trainer:log_validation_results:142 - CMC curve, Rank-5  :97.2%
2023-05-15 14:39:51.052 | INFO     | engine.trainer:log_validation_results:142 - CMC curve, Rank-10 :98.2%
​

知识蒸馏训练

支持网络为ResNet系列。

参数说明(基本参数和上面训练一样，只是多了kd)：

--model_name:模型名称，支持Resnet，resnet18_kd, resnet34_kd, resnet50_kd, resnet101_kd

--model_path:预权重路径

--kd:开启蒸馏模式

--feature_loss_coefficient:特征蒸馏的权重，默认0.03

这里用的蒸馏为在线式蒸馏(自蒸馏)，暂未更新离线式蒸馏。

python tools/train.py --model_name [model name] --model_path [your model weight path] --IMS_PER_BATCH 8 --TEST_IMS_PER_BATCH 4 --kd --feature_loss_coefficient 0.03

=> Market1501 loaded
Dataset statistics:
  ----------------------------------------
  subset   | # ids | # images | # cameras
  ----------------------------------------
  train    |   751 |    12936 |         6
  query    |   750 |     3368 |         6
  gallery  |   751 |    15913 |         6
  ----------------------------------------
resnet50_kd loading pretrained model weight...
label smooth on, numclasses: 751
ready kd train!
​

训练后会在logs文件下保存权重，命名格式为mAP_KD_xx.pth。下面是resnet50蒸馏前后第一个Epoch评价指标对比，还是有提升的【由于本人硬件环境受限，只是给大家把功能进行了实现】。

resnet50
Validation Results - Epoch: 1
2023-05-17 20:08:53.642 | INFO     | engine.trainer:log_validation_results:156 - mAP:39.2%
2023-05-17 20:08:53.642 | INFO     | engine.trainer:log_validation_results:158 - CMC curve, Rank-1  :65.6%
2023-05-17 20:08:53.642 | INFO     | engine.trainer:log_validation_results:158 - CMC curve, Rank-5  :80.3%
2023-05-17 20:08:53.642 | INFO     | engine.trainer:log_validation_results:158 - CMC curve, Rank-10 :85.0%
​
resnet50_kd:[layer3作为教师网络]
2023-05-17 20:22:07.030 | INFO     | engine.trainer:log_validation_results:153 - Validation Results - Epoch: 1
2023-05-17 20:22:07.131 | INFO     | engine.trainer:log_validation_results:156 - mAP:47.9%
2023-05-17 20:22:07.131 | INFO     | engine.trainer:log_validation_results:158 - CMC curve, Rank-1  :73.5%
2023-05-17 20:22:07.131 | INFO     | engine.trainer:log_validation_results:158 - CMC curve, Rank-5  :85.7%
2023-05-17 20:22:07.139 | INFO     | engine.trainer:log_validation_results:158 - CMC curve, Rank-10 :88.9%

教师网络&学生网络选择

教师网络采用深层网络，浅层网络为学生网络。

具体的教师网络和学生网络的选择可以看engine/trainer.py第46行至58行。

此处默认resnet中的layer3为教师网络，layer1，layer2为学生网络。具体的效果可以根据自己实际任务去尝试。采用特征蒸馏，暂未更新逻辑蒸馏。

       elif kd:
            score, feat, layer_out_feat = model(img)
            loss = loss_fn(score, feat, target)
            teacher_feature = layer_out_feat[1].detach()  # 取出教师层
            '''
            （rannge(idx,len(layer_out_feat))，中的idx可以决定哪个作为教师）
            idx=1表示layer4为教师网络,layer3,layer2,layer1为student
            idx=2表示layer3为教师网络,layer2,layer1为student
            idx=3表示layer2为教师网络，layer1为student
            '''
            for index in range(2, len(layer_out_feat)):  # layer4, layer3, layer2, layer1
                if index != 2:  # 排除自己
                    loss += torch.dist(layer_out_feat[index], teacher_feature) * feature_loss_coefficient

测试

python tools/test.py --TEST_IMS_PER_BATCH 4 --model_name [your model name] --model_path [your weight path]

可以进行mAP,Rank的测试

---

Reid相关资料学习链接

数据集代码详解：Reid训练代码之数据集处理_爱吃肉的鹏的博客-CSDN博客

Reid损失函数理论讲解：Reid之损失函数理论学习讲解_表征学习和度量学习_爱吃肉的鹏的博客-CSDN博客

Reid度量学习Triplet loss代码讲解：Reid度量学习Triplet loss代码解析。_爱吃肉的鹏的博客-CSDN博客

reid strong baseline代码学习：

Reid strong baseline 代码详解_爱吃肉的鹏的博客-CSDN博客 

预权重链接：

链接：百度网盘 请输入提取码 提取码：yypn

如果项目对你有用，麻烦点个Star

注：本项目暂时为免费开源，后期完善后会考虑适当收费【毕竟也是自己辛苦弄出来的】

项目代码

GitHub - YINYIPENG-EN/reid_strong_baseline_KD: reid strong baseline Knowledge distillation(reid知识蒸馏)

后期计划更新

1.引入知识蒸馏训练(已于2023.05.26更新)

2.加入YOLOX