这是一个讲解YOLO的入门教程的代码。在这个项目中,我们继承了YOLOv1~v4、YOLOX以及YOLOv7的中心思想,并在此基础上做了适当的修改来实现了结构较为简洁的YOLO检测器。我们希望通过初学者可以通过学习流行的YOLO检测器顺利入门目标检测领域。
| Model |
Backbone |
Scale |
IP |
Epoch |
APval 0.5
| FPS3090 FP32-bs1
| Weight |
| YOLOv1 |
ResNet-18 |
640 |
√ |
150 |
76.7 |
|
ckpt |
| YOLOv2 |
DarkNet-19 |
640 |
√ |
150 |
79.8 |
|
ckpt |
| YOLOv3 |
DarkNet-53 |
640 |
√ |
150 |
82.0 |
|
ckpt |
| YOLOv4 |
CSPDarkNet-53 |
640 |
√ |
150 |
83.6 |
|
ckpt |
| YOLOX |
CSPDarkNet-L |
640 |
√ |
150 |
84.6 |
|
ckpt |
| YOLOv7-Large |
ELANNet-Large |
640 |
√ |
150 |
86.0 |
|
ckpt |
所有的模型都使用了ImageNet预训练权重(IP),所有的FLOPs都是在VOC2007 test数据集上以640x640或1280x1280的输入尺寸来测试的。FPS指标是在一张3090型号的GPU上以batch size=1的输入来测试的,请注意,测速的内容包括模型前向推理、后处理以及NMS操作。
COCO
当COCO数据集的路径修改正确后,运行上述命令应该会看到COCO数据的可视化图像。
For example:
python train.py --cuda -d coco --root path/to/COCO -v yolov1 -bs 16 --max_epoch 150 --wp_epoch 1 --eval_epoch 10 --fp16 --ema --multi_scale
由于我的计算资源有限,我不得不在训练期间将batch size设置为16甚至更小。我发现,对于-Nano或-Tiny这样的小模型,它们的性能似乎对batch size不太敏感,比如我复制的YOLOv5-N和S,它们甚至比官方的YOLOv5-N和S略强。然而,对于*-Large这样的大模型,其性能明显低于官方的性能,这似乎表明大模型对batch size更敏感。
我提供了启用DDP训练的bash文件train_multi_gpus.sh,我希望有人可以使用更多的显卡和更大的batch size来训练我实现的大模型,如YOLOv5-L、YOLOX以及YOLOv7-L。如果使用更大的batch size所训练出来的性能更高,如果能将训练的模型分享给我,我会很感激的。
训练
使用单个GPU来训练
sh train_single_gpu.sh
使用者可以根据自己的情况来调整train_single_gpu.sh文件中的配置,以便在自己的本地上顺利训练模型。
如果使用者想查看训练时所使用的数据,可以在训练命令中输入--vsi_tgt参数,例如:
python train.py --cuda -d coco --root path/to/coco -v yolov1 --vis_tgt
使用多个GPU来训练
sh train_multi_gpus.sh
使用者可以根据自己的情况来调整train_multi_gpus.sh文件中的配置,以便在自己的本地上顺利训练模型。
当训练突然中断时, 使用者可以在训练命令中传入--resume参数,并指定最新保存的权重文件(默认为None),以便继续训练。例如:
python train.py \
--cuda \
-d coco \
-v yolov1 \
-bs 16 \
--max_epoch 150 \
--wp_epoch 1 \
--eval_epoch 10 \
--ema \
--fp16 \
--resume weights/coco/yolov1/yolov1_epoch_151_39.24.pth
测试
使用者可以参考下面的给出的例子在相应的数据集上去测试训练好的模型,正常情况下,使用者将会看到检测结果的可视化图像。
python test.py -d coco \
--cuda \
-v yolov1 \
--img_size 640 \
--weight path/to/weight \
--root path/to/dataset/ \
--show
对于YOLOv7,由于YOLOv7的PaFPN中包含了RepConv模块,因此你可以在测试命中加上--fuse_repconv来融合其中的RepConv:
python test.py -d coco \
--cuda \
-v yolov7_large \
--fuse_repconv \
--img_size 640 \
--weight path/to/weight \
--root path/to/dataset/ \
--show
验证
使用者可以参考下面的给出的例子在相应的数据集上去验证训练好的模型,正常情况下,使用者将会看到COCO风格的AP结果输出。
python eval.py -d coco-val \
--cuda \
-v yolov1 \
--img_size 640 \
--weight path/to/weight \
--root path/to/dataset/ \
--show
如果使用者想测试模型在COCO test-dev数据集上的AP指标,可以遵循以下步骤:
- 将上述命令中的
coco-val修改为coco-test,然后运行;
- 运行结束后,将会得到一个名为
coco_test-dev.json的文件;
- 将其压缩为一个
.zip,按照COCO官方的要求修改压缩文件的名称,例如``;
- 按照COCO官方的要求,将该文件上传至官方的服务器去计算AP。
Demo
本项目在data/demo/images/文件夹中提供了一些图片,使用者可以运行下面的命令来测试本地的图片:
python demo.py --mode image \
--path_to_img data/demo/images/ \
--cuda \
--img_size 640 \
-m yolov2 \
--weight path/to/weight \
--show
如果使用者想在本地的视频上去做测试,那么你需要将上述命令中的--mode image修改为--mode video,并给--path_to_vid传入视频所在的文件路径,例如:
python demo.py --mode video \
--path_to_img data/demo/videos/your_video \
--cuda \
--img_size 640 \
-m yolov2 \
--weight path/to/weight \
--show \
--gif
如果使用者想用本地的摄像头(如笔记本的摄像头)去做测试,那么你需要将上述命令中的--mode image修改为--mode camera,例如:
python demo.py --mode camera \
--cuda \
--img_size 640 \
-m yolov2 \
--weight path/to/weight \
--show \
--gif
检测的例子
运行命令如下:
python demo.py --mode video \
--path_to_vid ./dataset/demo/videos/000006.mp4 \
--cuda \
--img_size 640 \
-m yolov2 \
--weight path/to/weight \
--show \
--gif
结果如下:

目标跟踪
该项目也支持多目标跟踪任务。我们使用本项目的YOLO检测器作为“tracking-by-detection”的检测器,并使用简单高效的ByteTrack作为跟踪器。
images tracking
python track.py --mode image \
--path_to_img path/to/images/ \
--cuda \
-size 640 \
-dt yolov2 \
-tk byte_tracker \
--weight path/to/coco_pretrained/ \
--show \
--gif
video tracking
python track.py --mode video \
--path_to_img path/to/video/ \
--cuda \
-size 640 \
-dt yolov2 \
-tk byte_tracker \
--weight path/to/coco_pretrained/ \
--show \
--gif
camera tracking
python track.py --mode camera \
--cuda \
-size 640 \
-dt yolov2 \
-tk byte_tracker \
--weight path/to/coco_pretrained/ \
--show \
--gif
多目标跟踪的例子
- Detector: YOLOv2
- Tracker: ByteTracker
- Device: i5-12500H CPU
运行命令如下:
python track.py --mode video \
--path_to_img ./dataset/demo/videos/000006.mp4 \
-size 640 \
-dt yolov2 \
-tk byte_tracker \
--weight path/to/coco_pretrained/ \
--show \
--gif
结果如下:
