1.项目介绍

这次给大家介绍一个很方便的训练自己图像识别模型的一个程序。可以通过一行命令实现训练自己的图像识别模型,并且训练的速度很快,效果也不错。

图像分类有三种训练方式:

  • 构建一个新的模型并从头开始训练,称为scrach

  • 在已经训练好的模型基础上,修改模型的最后的全连接层,并重新训练全连接层称为bottleneck

  • 在已经训练好的模型基础上,修改模型的最后的全连接层,并重新训练全连接层同时微调模型的卷积层,称为finetune

我们这次项目的名称为bottleneck,所以我们要介绍的是第二种训练图像分类的方法。使用bottleneck的方式来训练自己的数据,优点是训练速度比较快,训练周期数比较少就可以得到比较好的结果。

这次训练模型使用的代码是tensorflow官方提供的图像重新训练的程序:
https://github.com/tensorflow/hub/tree/master/examples/image_retraining
该程序使用了一个已经训练好的inception-v3的模型来作为基础模型,inception-v3是使用imagenet比赛的数据集训练出来的一个可以识别1000种生活中常见物体的模型。

v3表示它的第三个版本,inception模型的使用是可以参考:
https://github.com/Qinbf/Tensorflow/tree/master/Tensorflow基础使用与图像识别应用/程序

项目所需环境:
Python
Tensorflow

2.Inception的bottleneck介绍

下图为inception模型的结构图:

如果把这个模型看成是一个花瓶的话,那么bottleneck瓶颈就是图中全连接FC(fully connected)的部分。

inception的模型本来是为了imagenet比赛而创建的,所以有1000个分类,全连接层有1000个输出。我们训练自己的数据集通常没有这么多分类,而且imagenet比赛中分类也不一定是我们想要的分类。

所以模型的全连接层的部分需要全部丢弃,然后创建新的全连接层并重新训练。比如我们要训练5个分类的图片,那么就创建一个有5个输出的全连接层来训练。

模型中的卷积层部分因为是经过大量训练得到的,所以已经具备了非常良好的特征提取的能力,所以就不需要重新训练了,可以不做修改直接使用。

训练流程:

  1. 先把数据集传入已经训练好的inception模型,得到FC层的前一层AveragePool层的输出,这个输出值相当于是从这张图片提取出的图像特征。然后把这个图像特征以文本形式保存到本地文件夹中。

  2. 根据图片类别数量重新构建全连接层,比如一共有5个类别的图片那么就会创建5个输出的全连接层。

  3. 把每张图片的特征输出传到新的全连接层中做训练,只训练全连接层。

  4. 把inception中的卷积层再加上后面训练好的全连接层就得到了新的模型,可以识别我们自己特定的数据集。

3.Inception模型准备

Inception-v3的模型比较大,所以我上传了百度云盘:

链接:百度云地址
密码:i5jw

在程序根目录创建一个inception_model的文件夹,然后把模型的压缩包下载到inception_model的文件夹中,然后解压。如图所示:

4.项目文件介绍

4.1 retrain.py为训练模型的程序

4.2 Inception_model为之前训练好的inception模型存放路径

4.3 data为需要训练的图片存放位置

图片存放的方式必须如上图所示,一个种类的图片放在一个文件夹下面。图片文件夹的名字就是图片类别的名字,名字要用英文。

4.4 bottleneck为inception卷积层提取出来的特征,以txt的格式保存在bottleneck文件夹中

4.5Linux和mac用户执行retrain.sh可以开始训练模型;windows用户执行retrain.bat可以开始训练模型)

retrain.bat内容如下:
python retrain.py ^
--bottleneck_dir bottleneck ^
--how_many_training_steps 200 ^
--model_dir inception_model ^
--output_graph output_graph.pb ^
--output_labels output_labels.txt ^
--image_dir data
pause

注:
bottleneck_dir :图片特征文件存放路径
how_many_training_steps :训练次数
model_dir:之前训练好的inception模型存放路径
output_graph:新训练的模型存放路径和名称
output_labels:生成的标签文件存放路径和名称
image_dir:训练数据存放路径

4.6 test_images为测试图片存放位置

4.7 predict.ipynb和predict.py为新训练好的模型的测试程序,执行这两个文件可以载入训练好的新的模型预测test_images文件夹中的图片类型。

5.训练和测试结果

我训练了5个种类的图片:'flower','guitar','house','animal','plane'

每个种类500张,特征提取的时间会比较长,我使用mac笔记本大概要二三十分钟,训练的时间很快,不到1分钟。

训练速度比较快的原因是训练次数我设置比较少只有200次,并且只训练了全连接层。训练结束后可以看到测试集准确率:

INFO:tensorflow:Final test accuracy = 99.6% (N=247)

几乎达到100%的准确率!

准确率之所以可以这么高主要原因有两个

第一是由于inception模型本身比较优秀,我们在一个优秀的模型上进行进一步训练所以得到的结果也很好;

第二是由于我使用的图像分类难度比较低,因为,花,吉他,房子,动物,飞机这几个种类之间的相似度非常低,所以很容易分辨。

如果是做狗的品种识别,比如狗有100个品种,因为狗与狗之间的相似度很高,那么准确率就不可能达到这么高了。

predict文件的部分测试效果

6.项目打包

大家可以去我的Github下载源代码:
qbf的Github