基于 TensorFlow 和 DALI 的 You Only Look Once v4

此示例介绍了一个 YOLOv4 网络的示例实现，该实现基于以下论文 – Alexey Bochkovskiy、Chien-Yao Wang、Hong-Yuan Mark Liao，YOLOv4：目标检测的最佳速度和准确率。

该实现基于 darknet 存储库中提供的代码。

要在所有可用的 GPU 上使用 DALI 从头开始运行训练，请运行以下命令

python src/main.py train /coco/train2017 /coco/annotations/instances_train2017.json \
  -b 8 -e 6 -s 1000 -o output.h5 \
  --pipeline dali-gpu --multigpu --use_mosaic

要每 epoch 保存检查点，请将以下标志添加到命令中

--ckpt_dir /ckpt

通过为 -w 标志提供适当的路径来恢复训练

-w /ckpt/epoch_3.h5

要每隔一个 epoch 执行评估，请添加以下标志

--eval_file_root /coco/val2017 \
--eval_annotations /coco/annotations/instances_val2017.json \
--eval_frequency 2 --eval_steps 500

要评估训练后的模型，请运行以下命令

python src/main.py eval /coco/val2017 /coco/annotations/instances_val2017.json \
  -w output.h5 -b 1 -s 5000

要执行推理并在屏幕上显示结果，请运行以下命令

python src/main.py infer image.png -w output.h5 -c coco-labels.txt

要求

• 此示例已使用 python 3.8 和 tensorflow 2.4.1 进行了测试，应该可以在更高版本上运行。

• 该模型需要 COCO 2017 数据集进行训练。

• 除了 TensorFlow 之外，还需要以下 python 包

  matplotlib tensorflow-addons pycocotools

用法

训练

usage: main.py train [-h] file_root annotations
  [--batch_size BATCH_SIZE] [--epochs EPOCHS] [--steps STEPS] [--output OUTPUT]
  [--start_weights START_WEIGHTS] [--log_dir LOG_DIR] [--ckpt_dir CKPT_DIR]
  [--pipeline PIPELINE] [--multigpu] [--use_mosaic] [--learning_rate LEARNING_RATE]
  [--eval_file_root EVAL_FILE_ROOT] [--eval_annotations EVAL_ANNOTATIONS]
  [--eval_steps EVAL_STEPS] [--eval_frequency EVAL_FREQUENCY]
  [--seed SEED]

参数	描述
-h, –help	显示此帮助信息并退出
file_root	包含 train coco2017 图像的文件夹路径
annotations	instances_train2017.json 文件路径
–batch_size BATCH_SIZE, -b BATCH_SIZE	每个训练步骤的图像数量，默认 = 8
–epochs EPOCHS, -e EPOCHS	训练 epoch 数量，默认 = 5
–steps STEPS, -s STEPS	每个 epoch 的训练步骤数，默认 = 1000
–output OUTPUT, -o OUTPUT	训练模型的 .h5 输出文件路径，默认 = output.h5
–start_weights START_WEIGHTS, -w START_WEIGHTS	h5 或 YOLO 格式的初始权重文件
–log_dir LOG_DIR	TensorBoard 日志目录路径
–ckpt_dir CKPT_DIR	检查点文件目录路径
–pipeline PIPELINE	dali_gpu、dali_cpu 或 numpy
–multigpu	如果存在，则使用所有可用 GPU 运行训练
–use_mosaic	如果存在，则使用 mosaic 数据增强
–learning_rate LEARNING_RATE	训练的学习率，默认 = 1e-3
–eval_file_root EVAL_FILE_ROOT	包含 val coco2017 图像的文件夹路径
–eval_annotations EVAL_ANNOTATIONS	instances_val2017.json 文件路径
–eval_steps EVAL_STEPS	每个评估步骤的图像数量，默认 = 5000
–eval_frequency EVAL_FREQUENCY	每次评估之间训练 epoch 的数量，默认 = 5
–seed SEED	DALI 和 TensorFlow 的种子

推理

usage: main.py infer [-h] image [--weights WEIGHTS] [--classes CLASSES]
                     [--output OUTPUT]

参数	描述
-h, –help	显示此帮助信息并退出
image	要在其上执行推理的图像路径
–weights WEIGHTS, -w WEIGHTS	h5 或 YOLO 格式的训练权重文件
–classes CLASSES, -c CLASSES	coco-labels.txt 文件路径
–output OUTPUT, -o OUTPUT	path to an output image

评估

usage: main.py eval [-h] file_root annotations [--weights WEIGHTS]
                    [--batch_size BATCH_SIZE] [--steps STEPS]

参数	描述
-h, –help	显示此帮助信息并退出
file_root	包含 val coco2017 图像的文件夹路径
annotations	instances_val2017.json 文件路径
–weights WEIGHTS, -w WEIGHTS	h5 或 YOLO 格式的训练权重文件
–batch_size BATCH_SIZE, -b BATCH_SIZE	每个评估步骤的图像数量，默认 = 1
–steps STEPS, -s STEPS	评估步骤数，默认 = 1000