以下是一个B站博主所写,讲的非常好,写的非常详细,这里在学习之余将该代码保存至git,针对代码的改动主要基于同门对数据接口的撰写,已使用于天池coco-json数据集,针对天池重庆工业酒瓶缺陷检测做了一个小demo,检测精度不会很高,仅作baseline用
- Python3.6或者3.7
- Pytorch1.6(注意:必须是1.6.0或以上,因为使用官方提供的混合精度训练1.6.0后才支持)
- pycocotools(Linux: pip install pycocotools;
Windows:pip install pycocotools-windows(不需要额外安装vs)) - Ubuntu或Centos(不建议Windows)
- 最好使用GPU训练
- 详细环境配置见
requirements.txt
* ├── backbone: 特征提取网络,可以根据自己的要求选择
* ├── network_files: Faster R-CNN网络(包括Fast R-CNN以及RPN等模块)
* ├── train_utils: 训练验证相关模块(包括cocotools)
* ├── my_dataset.py: 自定义dataset用于读取VOC数据集
* ├── train_mobilenet.py: 以MobileNetV2做为backbone进行训练
* ├── train_resnet50_fpn.py: 以resnet50+FPN做为backbone进行训练
* ├── train_multi_GPU.py: 针对使用多GPU的用户使用
* ├── predict.py: 简易的预测脚本,使用训练好的权重进行预测测试
* └── pascal_voc_classes.json: pascal_voc标签文件
- MobileNetV2 backbone: https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
- ResNet50+FPN backbone: https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
- Pascal VOC2012 train/val数据集下载地址:http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
- 如果不了解数据集或者想使用自己的数据集进行训练,请参考我的bilibili:https://b23.tv/F1kSCK
- 确保提前准备好数据集
- 确保提前下载好对应预训练模型权重
- 若要训练mobilenetv2+fasterrcnn,直接使用train_mobilenet.py训练脚本
- 若要训练resnet50+fpn+fasterrcnn,直接使用train_resnet50_fpn.py训练脚本
- 若要使用多GPU训练,使用
python -m torch.distributed.launch --nproc_per_node=8 --use_env train_multi_GPU.py
指令,nproc_per_node
参数为使用GPU数量