ImageAI训练自定义数据总结

1.先按《Win10下ImageAI-gpu(tensorflow-gpu)环境安装（cuda及cuDNN自动安装）及实时对象检测》安装好。

给图片进行对象标注使用LabelImage工具。

1. 下载pretrained-yolov3.h5，放到temp.py同一目录下。

准备自己的数据集（我这里文件夹是net）放到temp.py同一目录下。

数据集目录结构必须包含train和validation

train文件夹下必须包含images和annotations.其中image用于存放图片，annotations用于存放xml

 

validation文件夹结构与train结构相同

 

1. 编写代码开始训练。

设置数据集文件夹，这里为net

trainer.setDataDirectory(data_directory="net")

设置训练配置信息

trainer.setTrainConfig(object_names_array=["net","pillar"], batch_size=2, num_experiments=100,

                       train_from_pretrained_model="pretrained-yolov3.h5")

其中object_names_array=["net","pillar"]，设置训练对象，必须跟xml对象名称一致,否则会报错

可以打开annotations下面的xml文件查看需要训练的对象

batch_size=2是关键，否则可能出现以下错误

 

Resource exhausted: OOM when allocating tensor with shape[8,1024,14,14] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc

完整代码如下：

from imageai.Detection.Custom import DetectionModelTrainer

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

 

1. 开始训练

正在训练