使用TensorFlow 2.0实现的YoloV3
这个仓库提供了使用TensorFlow 2.0实现的干净的YoloV3实现方式,遵循所有最佳实践。
关键特性
- TensorFlow 2.0
- 带有预训练权重的
yolov3
- 带有预训练权重的
yolov3-tiny
- 推理示例
- 迁移学习示例
- 使用
tf.GradientTape
的 eager 模式训练 - 使用
model.fit
的图模式训练 - 使用
tf.keras.layers
的功能模型 - 使用
tf.data
的输入管道 - Tensorflow Serving
- 向量化转换
- GPU 加速
- 完全整合
absl-py
from abseil.io - 干净的实现
- 遵循最佳实践
- MIT 许可证
使用方法
安装
Conda(推荐)
# Tensorflow CPU
conda env create -f conda-cpu.yml
conda activate yolov3-tf2-cpu
# Tensorflow GPU
conda env create -f conda-gpu.yml
conda activate yolov3-tf2-gpu
Pip
pip install -r requirements.txt
Nvidia 驱动(用于 GPU)
# Ubuntu 18.04
sudo apt-add-repository -r ppa:graphics-drivers/ppa
sudo apt install nvidia-driver-430
# Windows/其他
https://www.nvidia.com/Download/index.aspx
转换预训练 Darknet 权重
# 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
# 摄像头
python detect_video.py --video 0
# 视频文件
python detect_video.py --video path_to_file.mp4 --weights ./checkpoints/yolov3-tiny.tf --tiny
# 带输出的视频文件
python detect_video.py --video path_to_file.mp4 --output ./output.avi
训练
我已经创建了一个完整的教程,说明如何使用VOC2012数据集从头开始训练。 查看文档:https://github.com/zzh8829/yolov3-tf2/blob/master/docs/training_voc.md
对于自定义训练,您需要按照TensorFlow物体检测API生成tfrecord。 例如,您可以使用Microsoft VOTT生成这样的数据集。 您还可以使用这个脚本创建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
python train.py --batch_size 8 --dataset ~/Data/voc2012.tfrecord --val_dataset ~/Data/voc2012_val.tfrecord --epochs 100 --mode fit --transfer none
python train.py --batch_size 8 --dataset ~/Data/voc2012.tfrecord --val_dataset ~/Data/voc2012_val.tfrecord --epochs 100 --mode fit --transfer no_output
python train.py --batch_size 8 --dataset ~/Data/voc2012.tfrecord --val_dataset ~/Data/voc2012_val.tfrecord --epochs 10 --mode eager_fit --transfer fine_tune --weights ./checkpoints/yolov3-tiny.tf --tiny
TensorFlow 服务
你可以将模型导出到 tf 服务
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_iamges
)
输出是
yolo_nms_0:边界框
yolo_nms_1:分数
yolo_nms_2:类别
yolo_nms_3:有效检测数
基准测试(尚未训练)
通过 detect_video.py
粗略计算得出的数字
Macbook Pro 13 (2.7GHz i5)
检测 | 416x416 | 320x320 | 608x608 |
---|---|---|---|
YoloV3 | 1000ms | 500ms | 1546ms |
YoloV3-Tiny | 100ms | 58ms | 208ms |
台式机 PC (GTX 970)
检测 | 416x416 | 320x320 | 608x608 |
---|---|---|---|
YoloV3 | 74ms | 57ms | 129ms |
YoloV3-Tiny | 18ms | 15ms | 28ms |
AWS g3.4xlarge (Tesla M60)
检测 | 416x416 | 320x320 | 608x608 |
---|---|---|---|
YoloV3 | 66ms | 50ms | 123ms |
YoloV3-Tiny | 15ms | 10ms | 24ms |
RTX 2070 (感谢 @AnaRhisT94)
检测 | 416x416 |
---|---|
使用 predict_on_batch 的 YoloV3 | 29-32ms |
使用 predict_on_batch + TensorRT 的 YoloV3 | 22-28ms |
在 416x416 分辨率下,使用 Titan X 的 Darknet 版本 YoloV3 需时 29ms。 考虑到 Titan X 的评价约为 Tesla M60 的两倍, 从性能角度来看,这个实现是相当不错的。
实现细节
Eager 执行
对于现有 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
对于调试非常有用,你可以在任何地方设置断点。
你可以使用 model.compile 中的 run_eagerly
参数,
结合所有 keras 的拟合功能与 gradient tape。
我有限的测试中,所有的训练方法,包括 GradientTape,keras.fit,
无论是 eager 还是非 eager 模式,性能相似。但图模式仍然优选,
因为它效率更高。
@tf.function
@tf.function 非常酷。它像是 eager 与图之间的中间版本。 你可以通过禁用 tf.function 来逐步调试该函数, 然后在生产中启用它以获得性能提升。重要注意事项是, 不要传递任何非 tensor 参数给 @tf.function,这将导致每次调用时重新编译。 我不确定除了使用全局变量之外最好的方法是什么。
absl.py (abseil)
绝对惊人。如果你还不知道,absl.py 已被谷歌内部项目正式使用。它标准化了 Python 和许多其他语言的应用程序接口。在谷歌内部使用之后,我很高兴听到 abseil 开源。它包含了几十年来从创建大型可扩展应用程序中学到的最佳实践。我真的没有什么不好的话要说,强烈推荐 absl.py 给大家。
加载预训练的 Darknet 权重
使用纯函数式 API 非常困难,因为 tf.keras 和 darknet 中的层顺序不同。这里的干净解决方案是在 keras 中创建子模型。Keras 无法正确地以 h5 格式保存嵌套模型,建议使用 TF Checkpoint,因为它得到了 TensorFlow 的官方支持。
tf.keras.layers.BatchNormalization
它在迁移学习中效果不佳。互联网上有很多文章和 github 问题。我用一个简单的黑客方法让它在小批量迁移学习中更好地工作。
transform_targets 的输出是什么???
我知道这很令人困惑,但输出是形状为
(
[N, 13, 13, 3, 6],
[N, 26, 26, 3, 6],
[N, 52, 52, 3, 6]
)
的元组,其中 N 是批次中的标签数,最后一个维度 "6" 代表边界框的 [x, y, w, h, obj, class]
。
IOU 和分数阈值
默认阈值是 IOU 和分数的 0.5,您可以根据需要通过设置 --yolo_iou_threshold
和 --yolo_score_threshold
标志来调整它们。
最大框数量
默认情况下,每张图像最多可以有 100 个边界框,如果出于某种原因您想要更多的框,可以使用 --yolo_max_boxes
标志。
NAN 损失 / 训练失败 / 不收敛
包括我在内的许多人都成功进行了训练,所以代码肯定有效 @LongxingTan 在 https://github.com/zzh8829/yolov3-tf2/issues/128 中提供了一些他的见解,摘要如下:
- 对于 nan 损失,尝试将学习率调小
- 仔细检查您的输入数据格式。通过 vott 和 labelImg 标记的数据输入是不同的,所以请确保输入框是正确的,并仔细检查格式是
x1/width,y1/height,x2/width,y2/height
而 不是 x1,y1,x2,y2,或 x,y,w,h
确保使用此工具可视化您的自定义数据集
python tools/visualize_dataset.py --classes=./data/voc2012.names
它将输出数据集中随机图像到 output.jpg
如果渲染的标签看起来不正确,培训绝对不会有效
命令行参数参考
convert.py:
--output: 输出路径
(默认: './checkpoints/yolov3.tf')
--[no]tiny: yolov3 或 yolov3-tiny
(默认: 'false')
--weights: 权重文件路径
(默认: './data/yolov3.weights')
--num_classes: 模型中的类数
(默认: '80')
(一个整数)
detect.py:
--classes: 类文件路径
(默认: './data/coco.names')
--image: 输入图像路径
(默认: './data/girl.png')
--output: 输出图像路径
(默认: './output.jpg')
--[no]tiny: yolov3 或 yolov3-tiny
(默认: 'false')
--weights: 权重文件路径
(默认: './checkpoints/yolov3.tf')
--num_classes: 模型中的类数
(默认: '80')
(一个整数)
detect_video.py:
--classes: 类文件路径
(默认: './data/coco.names')
--video: 输入视频路径 (使用 0 表示摄像头)
(默认: './data/video.mp4')
--output: 输出视频路径(记得为给定格式设置正确的编码。例如:XVID 对于 .avi)
(默认: 无)
--output_format: 保存视频到文件时使用的编码
(默认: 'XVID')
--[no]tiny: yolov3 或 yolov3-tiny
(默认: 'false')
--weights: 权重文件路径
(默认: './checkpoints/yolov3.tf')
--num_classes: 模型中的类数
(默认: '80')
(一个整数)
train.py:
--batch_size: 批处理大小
(默认: '8')
(一个整数)
--classes: 类文件路径
(默认: './data/coco.names')
--dataset: 数据集路径
(默认: '')
--epochs: 纪元数
(默认: '2')
(一个整数)
--learning_rate: 学习率
(默认: '0.001')
(一个数字)
--mode: <fit|eager_fit|eager_tf>: fit: model.fit, eager_fit: model.fit(run_eagerly=True), eager_tf: 自定义 GradientTape
(默认: 'fit')
--num_classes: 模型中的类数
(默认: '80')
(一个整数)
--size: 图像大小
(默认: '416')
(一个整数)
--[no]tiny: yolov3 或 yolov3-tiny
(默认: 'false')
--transfer: <none|darknet|no_output|frozen|fine_tune>: none: 从头训练, darknet: 迁移 darknet, no_output: 迁移除输出外的所有内容, frozen: 迁移并冻结所有内容, fine_tune: 迁移所有并只冻结 darknet
(默认: 'none')
--val_dataset: 验证数据集路径
(默认: '')
--weights: 权重文件路径
(默认: './checkpoints/yolov3.tf')
更新日志
2019 年 10 月 1 日
- 更新至 Tensorflow v2.0.0 版本
参考资料
仅从 yolov3 论文中几乎不可能实现这一点。我不得不参考官方(非常难以理解)和许多非官方(有很多小错误)的 repo 来拼凑出完整的图景。