训练目标检测模型只需要这 6 行代码

AI科技评论

发布于 2019-08-21 15:19:18

1.4K0

发布于 2019-08-21 15:19:18

文章被收录于专栏：AI科技评论AI科技评论

在你的定制数据集上训练目标检测模型的分步教程。

原标题 | Train Object Detection AI with 6 lines of code

作者 | Moses Olafenwa

翻译 | 珺毅（浙江师范大学）

编辑 | Pita

在你的定制数据集上训练目标检测模型的分步教程。

目标检测是计算机视觉意义最深远的层面之一，因为它要求你在图片或者视频中定位、识别、跟踪感兴趣的物体。目标检测在一些有意思的工作或者研究领域被广泛的应用，例如：

自动驾驶
安全方面
行人/人群检测
车牌号车辆检测
工业自动化（例如，物品的拣选和分拣）
机器人及更多

许多预先收集的对象检测数据集，如Pascal VOC、Microsoft的COCO、谷歌的Open Images，以及它们的预先训练的模型，都可以很容易地用于检测和识别一组固定的项。

然而，使用这些现有的公共数据集和已经预训练好的模型所面临的挑战是，它们不会给你提供一个非常方便简洁的方式让你来训练一个新的目标检测模型，让你来检测和识别出自己理想中感兴趣的物体。从去年我出版了我的第一篇目标检测的文章《10行代码的目标检测》，我已经收到了来自开发者，项目团队，学生还有研究员们的数千条请求，他们想在提供了超过了80种类别的COCO数据集和提供了超过了600种类别的Open Image数据集中的图片和视频中检测出他们自己感兴趣的物体。

我非常欣慰的宣布伴随着ImageAI v2.1.0版本的发布，它可以完全支持训练你定制的YOLOv3模型去检测字面上任何种类和数量的对象，这就是我在这个教程中所要指导你所做的，让我们开始今天的课程吧。

OlafenwaMoses/ImageAI（https://github.com/OlafenwaMoses/ImageAI）

出于这个教程的目的，我们将使用Google Colab来训练我们已经提供的一个样本数据集，步骤如下。

步骤一 —— 准备你的数据集

对于自定义的检测训练，必须提供样本图片（你的图片数据集）来训练你的模型并且在精确度训练以后来验证这个模型。 ImageAI检测训练支持Pascal VOC格式的自定义数据集。就本教程而言，我们提供了Hololens混合现实耳机的样本数据集，在这个数据集上我们将要训练一个能够在图片和视频中检测和识别出Hololens的模型。你可以通过下面的链接下载这个样本数据集。

https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/headset.zip

如果你想在你的图片集上训练自己自定义的物体，请按照下面的指示：

决定你所要检测物体的类型，搜集关于这个物体200个左右（最低限度）或者更多的图片。
一旦你搜集到了这些图片，你需要在图片中注释这个物体。ImageAI使用Pascal VOC格式在图片进行了注释。你可以按照我们通过链接提供的分布教程为你的图像生成此注释。

参考文章：目标检测训练——准备你自定义的数据集（https://medium.com/deepquestai/object-detection-training-preparing-your-custom-dataset-6248679f0d1d）

一旦对你所有的照片增加了注释，给你的数据集建立一个文件夹并且在这个母文件夹中，在创建一个子文件夹和验证文件夹。
在训练文件夹中，创建图片和和注释子文件夹。把大约70%-80%数据图片放到图片文件夹中，并且把这些图片相应的注释放到注释文件夹中。
一旦你完成了上述的所有的工作，你的图片数据集文件夹应该跟下面的样本类似：

步骤二 —— 安装ImageAI和Dependencies

访问 https://colab.research.google.com 这个网站并且创建一个新的python3笔记本，确保将新笔记本运行时更改为GPU状态，然后：

运行下面的命令来下载hololens样本数据集

!wget 
https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/hololens.zip

解压缩hololens数据集的zip文件

!unzip hololens.zip

出于训练的目的，安装 Tensorflow GPU版本1.13.1。这是因为安装在你新笔记本上默认的Tensorflow GPU版本，当你训练你的模型的时候，会生成‘_TfDeviceCaptureOp’ 错误，物体目标没有这个属性‘_set_device_from_string’。如果你的新笔记本在安装以后收到一个要重新启动的提示，那就按照指示重新启动。

!pip3 install tensorflow-gpu==1.13.1

然后从PIP安装最近版本的ImageAI

!pip3 install imageai --upgrade

步骤三 —— 初始化你的检测训练模型

为了确保我们自定义的模型有更好的检测精确度，我们将在训练过程阶段从预训练的YOLOv3模型中使用迁移学习。ImageAI提供了使用和不使用迁移学习的选项。在这里我强烈建议你使用迁移学习，除非在你数据集里有成千上万个样本。

下载预训练的YOLOv3模型

!wget 
https://github.com/OlafenwaMoses/ImageAI/releases/download/essential-v4/pretrained-yolov3.h5

如下所示运行训练代码

from imageai.Detection.Custom import DetectionModelTrainer


trainer = DetectionModelTrainer()    
trainer.setModelTypeAsYOLOv3()    
trainer.setDataDirectory(data_directory="hololens")    
trainer.setTrainConfig(object_names_array=["hololens"], batch_size=4, num_experiments=100, train_from_pretrained_model="pretrained-yolov3.h5")    
trainer.trainModel()

简单理解一下！以上6行代码是你必须要做的来在你的自定义数据集上初始化训练。现在我们将以上代码拆解为下面的部分：

在第一行，我们从ImageAI导入‘DetectionModelTrainer’类
在第二行和第三行，我们创建了一个类的实例，并且把我们的模型的类型设置为YOLOv3
在第四行，我们将路径设置为我们自定义数据集的路径
在第五行，我们指定下面的参数：
- object_names_array：这是你的数据集中所有类别名称的一个数组。
- batc_size：这是训练的批次大小。请注意，批次大小越大，对于以保存的模型检测精确度越高。然而，鉴于在Colab上英伟达显卡K80的内存限制，我们必须把这个数值设置为4。
- num_experiments：这是我们希望训练代码在自定义数据集上的迭代次数。
- train_from_pretrained_model：这利用我们之前下载的预训练YOLOv3模型来利用迁移学习。

一旦训练开始，

ImageAI将会在hololens/jsonfolder中产生detection_config.json文件。这个JSON文件在图片或者视频的目标检测中使用。
ImageAI将会创建hololens/models文件夹，这里将存储所有产生的模型。
你将会在日志中看到如下示例的详细信息。

Using TensorFlow backend.
Generating anchor boxes for training images and annotation...
Average IOU for 9 anchors: 0.78
Anchor Boxes generated.
Detection configuration saved in  hololens/json/detection_config.json
Training on:   ['hololens']
Training with Batch Size:  4
Number of Experiments:  200


Epoch 1/200
 - 733s - loss: 34.8253 - yolo_layer_1_loss: 6.0920 - yolo_layer_2_loss: 11.1064 - yolo_layer_3_loss: 17.6269 - val_loss: 20.5028 - val_yolo_layer_1_loss: 4.0171 - val_yolo_layer_2_loss: 7.5175 - val_yolo_layer_3_loss: 8.9683
Epoch 2/200
 - 648s - loss: 11.1396 - yolo_layer_1_loss: 2.1209 - yolo_layer_2_loss: 4.0063 - yolo_layer_3_loss: 5.0124 - val_loss: 7.6188 - val_yolo_layer_1_loss: 1.8513 - val_yolo_layer_2_loss: 2.2446 - val_yolo_layer_3_loss: 3.5229
Epoch 3/200
 - 674s - loss: 6.4360 - yolo_layer_1_loss: 1.3500 - yolo_layer_2_loss: 2.2343 - yolo_layer_3_loss: 2.8518 - val_loss: 7.2326 - val_yolo_layer_1_loss: 1.8762 - val_yolo_layer_2_loss: 2.3802 - val_yolo_layer_3_loss: 2.9762
Epoch 4/200
 - 634s - loss: 5.3801 - yolo_layer_1_loss: 1.0323 - yolo_layer_2_loss: 1.7854 - yolo_layer_3_loss: 2.5624 - val_loss: 6.3730 - val_yolo_layer_1_loss: 1.4272 - val_yolo_layer_2_loss: 2.0534 - val_yolo_layer_3_loss: 2.8924
Epoch 5/200
 - 645s - loss: 5.2569 - yolo_layer_1_loss: 0.9953 - yolo_layer_2_loss: 1.8611 - yolo_layer_3_loss: 2.4005 - val_loss: 6.0458 - val_yolo_layer_1_loss: 1.7037 - val_yolo_layer_2_loss: 1.9754 - val_yolo_layer_3_loss: 2.3667
Epoch 6/200
 - 655s - loss: 4.7582 - yolo_layer_1_loss: 0.9959 - yolo_layer_2_loss: 1.5986 - yolo_layer_3_loss: 2.1637 - val_loss: 5.8313 - val_yolo_layer_1_loss: 1.1880 - val_yolo_layer_2_loss: 1.9962 - val_yolo_layer_3_loss: 2.6471
Epoch 7/200

步骤四 —— 评估你的模型

在如上所示的样本日志中，新的样本将会基于验证损失（例如，损失为5.2569）的下降而被存储。在大多数情况下，损失越低，模型在视频或者图片中检测物体的精度也就越高。然而，一些模型可能会出现过拟合并且仍然保持着低损失值。为了确保选出对于你自定义检测最佳的模型，ImageAI允许你评估存储在hololens/models文件夹中已训练模型的mAP值（平均精度，点击这里了解更多）。

mAP的值越高，模型的检测精度就越高。

在训练期间保存的模型上简单运行一下代码：

from imageai.Detection.Custom import DetectionModelTrainer

trainer = DetectionModelTrainer()
trainer.setModelTypeAsYOLOv3()
trainer.setDataDirectory(data_directory="hololens")
trainer.evaluateModel(model_path="hololens/models", json_path="hololens/json/detection_config.json", iou_threshold=0.5, object_threshold=0.3, nms_threshold=0.5)

运行如上代码时，你就会得到如下例子的结果。

Model File:  hololens/models/detection_model-ex-07--loss-4.42.h5
Using IoU :  0.5
Using Object Threshold :  0.3
Using Non-Maximum Suppression :  0.5
hololens: 0.9231
mAP: 0.9231
===============================
Model File:  hololens/models/detection_model-ex-10--loss-3.95.h5
Using IoU :  0.5
Using Object Threshold :  0.3
Using Non-Maximum Suppression :  0.5
hololens: 0.9725
mAP: 0.9725
===============================
Model File:  hololens/models/detection_model-ex-05--loss-5.26.h5
Using IoU :  0.5
Using Object Threshold :  0.3
Using Non-Maximum Suppression :  0.5
hololens: 0.9204
mAP: 0.9204
===============================

让我们拆分一下评估代码：

在前四行，我们导入了相同的训练类，创建了类实例，设置了检测模型的类型并且把路径设置为我们数据集的根目录。
在第五行，我们访问了评估模型函数并且在下面指定说明了参数。
- model_path：这是文件夹的目录且包含我们的模型，它也是对于一个指定模型的路径。
- json_file：这是在训练期间保存detection_config.json的文件夹。
- iou_threshold：这是我们所希望对于mAP计算的最低限度的IoU。它的值被设置在0.0与1.0之间。
- object_threshold：这是我们所希望对于mAP计算的最小类得分。它的值被设置在0.0与1.0之间。
- nms_threshold：这是我们所希望对于mAP计算的非极大抑制值。

步骤五 —— 在一张图片内检测自定义的物体

现在我们已经训练好了我们的自定义模型来检测Hololens headset，我们将使用已存储的最好的模型还有生成的detection_config.json文件来检测一张图片的物体。