在本系列文章中,探讨了如何在Android应用中使用TensorFlow Lite进行深度学习。从设置开发环境开始,寻找预训练模型,并将其适配到Android应用中。在前几部分,使用YOLO实现创建了一个能够识别80类对象的应用。本文将介绍如何训练自己的YOLO网络,以识别选择的项目分类。
YOLO(You Only Look Once)最初由Joseph Redmond和Ali Farhadi使用Darknet实现。Darknet是一个基于C的开源神经网络框架。如果机器已经安装了C编译器和make工具,那么可以通过以下三个命令下载并构建Darknet:
git clone https://github.com/pjreddie/darknet
cd darknet
make
确定希望网络识别的图像类别。在文本文件中列出这些类别,以便它们按照定义的顺序排列。每个类别应该单独一行。这个文件将用于训练。有序列表很重要,因为图像类别将通过这个列表中的零基索引进行注释。将此文件命名为classes.txt,并将其放置在Darknet的data文件夹中。
收集不同类别的对象图像进行训练和测试可能是最具挑战性的步骤之一,因为每个类别需要数百张图像。对于想要识别的物品,需要从不同角度和不同情境下的图片。每个图像文件还必须被标记。通过创建一个与图像同名的文本文件(但扩展名为txt)来标记图像,其中包含训练图像中每个样本对象的一行。每行包含五个由空格分隔的元素:
例如,以下行表示一个被归类为类别5的对象位于图像的x,y位置10,12,宽度为45像素,高度为22像素。由于使用了零基索引,这里的数字5代表之前创建的classes.txt文件中的第六行。
5 10 12 45 22
有了标记的图像后,需要将图像分成两个列表。这些列表应该按照80%/20%或90%/10%的比例划分。较大的列表将用于训练网络,较小的列表将用于测试网络。每个列表应该包含图像文件的绝对路径,每个路径单独一行。将这些文件命名为test.txt和train.txt。
现在,将这些文件作为配置集的一部分。在data文件夹中创建一个名为training_config.data的文件。该文件将包含网络将检测的类别数量、包含训练数据文件列表的路径、包含测试文件列表的路径,以及在训练过程中放置中间结果的文件夹名称。
classes=20
train=/dev/darknet/data/train.txt
valid=/dev/darknet/data/test.txt
names=data/classes.txt
backup=backup
有了这个文件,现在可以开始训练了:
./darknet detector train cfg/training_config.data cfg/yolov3-voc.cfg
训练是一个耗时的过程。一旦训练开始,可能需要离开电脑去做其他事情。如果在笔记本电脑上运行训练,请确保它连接到电源。训练计算密集型,电脑可能会长时间处于较高的功耗状态。
训练完成后,将得到一个DarkNet.weights文件。有几个开源工具可以将此文件转换为TensorFlow冻结图,用于TensorFlow 1或TensorFlow 2。建议使用的TensorFlow兼容工具名为Yolov4-tflite,这是一个Python脚本工具,可以从以下URL克隆:
https://github.com/hunglc007/tensorflow-yolov4-tflite
要转换权重文件,该工具接受文件的路径:
python save_model.py --weights ../darknet/yolov3.weights -output my-yolov --input_size 416 --model yolov3
程序的输出提供了一个TensorFlow冻结图,准备使用或转换为TensorFlow Lite。