YOLOv3在TensorFlow 2.0中的实现:功能强大的目标检测模型

YOLOv3在TensorFlow 2.0中的实现

YOLOv3是一种快速、准确的目标检测算法,在计算机视觉领域广受欢迎。本文将详细介绍YOLOv3在TensorFlow 2.0中的实现,包括其主要特性、安装使用方法、训练检测过程以及实现细节等。

主要特性

这个YOLOv3-TF2实现具有以下主要特性:

基于TensorFlow 2.0
支持yolov3和yolov3-tiny预训练权重
提供推理和迁移学习示例
支持eager模式和graph模式训练
使用tf.keras.layers构建函数式模型
使用tf.data构建输入pipeline
支持TensorFlow Serving
GPU加速
与absl-py完全集成
代码实现清晰,遵循最佳实践

安装与使用

安装

推荐使用Conda安装:

# CPU版本
conda env create -f conda-cpu.yml
conda activate yolov3-tf2-cpu

# GPU版本  
conda env create -f conda-gpu.yml
conda activate yolov3-tf2-gpu

也可以使用pip安装:

pip install -r requirements.txt

转换预训练权重

# yolov3
wget https://pjreddie.com/media/files/yolov3.weights -O data/yolov3.weights
python convert.py --weights ./data/yolov3.weights --output ./checkpoints/yolov3.tf

# yolov3-tiny
wget https://pjreddie.com/media/files/yolov3-tiny.weights -O data/yolov3-tiny.weights  
python convert.py --weights ./data/yolov3-tiny.weights --output ./checkpoints/yolov3-tiny.tf --tiny

目标检测

# yolov3
python detect.py --image ./data/meme.jpg

# yolov3-tiny
python detect.py --weights ./checkpoints/yolov3-tiny.tf --tiny --image ./data/street.jpg

YOLOv3检测结果

训练

作者提供了一个使用VOC2012数据集从头开始训练的完整教程,可以参考这里。

对于自定义数据集,需要生成tfrecord格式。可以使用Microsoft VOTT等工具生成数据集,也可以使用TensorFlow Object Detection API提供的脚本创建Pascal VOC格式的数据集。

训练命令示例:

python train.py --batch_size 8 \
                --dataset ~/Data/voc2012.tfrecord \
                --val_dataset ~/Data/voc2012_val.tfrecord \
                --epochs 100 \
                --mode eager_tf \
                --transfer fine_tune

实现细节

Eager execution

Eager execution对于现有的TensorFlow专家来说是一个很好的补充。它不太容易使用,需要对TensorFlow图有一定的中等理解。当你意外使用了不兼容的特性(如tensor.shape[0])或者在eager模式下可以正常工作但在尝试将模型编译为图时完全崩溃的Python控制流时,会很烦人。

model(x) vs. model.predict(x)

直接调用model(x)时,我们是在eager模式下执行图。对于model.predict,tf实际上是在第一次运行时编译图,然后在图模式下执行。因此,如果你只运行模型一次,model(x)会更快,因为不需要编译。否则,model.predict或使用导出的SavedModel图会快得多(快2倍)。对于非实时使用,model.predict_on_batch甚至更快(由@AnaRhisT94测试)。

GradientTape

GradientTape对于调试非常有用,你可以在任何地方设置断点。你可以使用model.compile中的run_eagerly参数将所有keras拟合功能与梯度带编译在一起。从我有限的测试来看,包括GradientTape、keras.fit在内的所有训练方法,无论是否eager,都产生了类似的性能。但图模式仍然是首选,因为它稍微更高效一些。

@tf.function

@tf.function非常酷。它就像eager和graph之间的中间版本。你可以通过禁用tf.function来逐步执行函数,然后在生产中启用它以获得性能。重要的是,你不应该向@tf.function传递任何非张量参数,它会导致每次调用时重新编译。我不确定除了使用全局变量之外还有什么最好的方法。

absl.py (abseil)

绝对令人惊叹。如果你还不知道,absl.py在Google内部项目中正式使用。它为Python和许多其他语言标准化了应用程序接口。在Google内部使用它之后,我很高兴听到abseil开源。它包含了从创建大规模可扩展应用程序中学到的几十年最佳实践。我真的没有什么坏话要说,强烈推荐每个人使用absl.py。

TensorFlow Serving

你可以将模型导出到TensorFlow Serving:

python export_tfserving.py --output serving/yolov3/1/
# 验证tfserving图
saved_model_cli show --dir serving/yolov3/1/ --tag_set serve --signature_def serving_default

输入是预处理的图像(参见dataset.transform_images),输出包括:

yolo_nms_0: 边界框
yolo_nms_1: 分数
yolo_nms_2: 类别
yolo_nms_3: 有效检测数量

性能基准

作者提供了在不同硬件上的性能基准测试结果,包括MacBook Pro、桌面PC(GTX 970)、AWS g3.4xlarge(Tesla M60)等。结果显示该实现的性能与Darknet版本的YOLOv3相当。

常见问题

NAN Loss / 训练失败 / 不收敛

许多人(包括作者)都成功进行了训练,所以代码肯定是可行的。@LongxingTan在#128中提供了一些见解:

对于nan损失,尝试减小学习率
仔细检查输入数据的格式,确保边界框格式正确

可以使用以下工具可视化自定义数据集:

python tools/visualize_dataset.py --classes=./data/voc2012.names

它会将数据集中的一个随机图像及其标签输出到output.jpg。如果渲染的标签看起来不正确,训练肯定无法正常工作。

结论

YOLOv3-TF2是一个功能强大、性能优异的目标检测模型实现。它充分利用了TensorFlow 2.0的新特性,同时保持了良好的代码结构和实践。无论是用于研究还是实际应用,YOLOv3-TF2都是一个值得尝试的选择。希望本文能帮助读者更好地了解和使用YOLOv3-TF2,在计算机视觉任务中取得更好的效果。

YOLOv3-TF2检测结果