使用自己的数据集训练MobileNet、ResNet实现图像分类（TensorFlow）| CSDN博文精选

作者 | pan_jinquan

来源 | CSDN博文精选

之前写了一篇博客《使用自己的数据集训练GoogLenet InceptionNet V1 V2 V3模型（TensorFlow）》https://panjinquan.blog.csdn.net/article/details/81560537，本博客就是此博客的框架基础上，完成对MobileNet的图像分类模型的训练，其相关项目的代码也会统一更新到一个Github中，强烈建议先看这篇博客《使用自己的数据集训练GoogLenet InceptionNet V1 V2 V3模型（TensorFlow）》后，再来看这篇博客。

TensorFlow官网中使用高级API -slim实现了很多常用的模型，如VGG，GoogLenet V1、V2和V3以及MobileNet、resnet模型，可详看这里https://github.com/tensorflow/models/tree/master/research/slim，当然TensorFlow官网也提供了训练这些模型的脚本文件，但灵活性太差了，要想增加log或者其他信息，真的很麻烦。

本人花了很多时间，去搭建一个较为通用的模型训练框架《tensorflow_models_nets》，目前几乎可以支持所有模型的训练，由于训练过程是自己构建的，所以你可以在此基础上进行任意的修改，也可以搭建自己的训练模型。

重要说明：
（1）项目Github源码：https://github.com/PanJinquan/tensorflow_models_learning，麻烦给个“Star”
（2）你需要一台显卡不错的服务器，不然会卡的一比，慢到花都谢了
（3）对于MobileNet、resnet等大型的网络模型，重头开始训练，是很难收敛的。但迁移学习finetune部分我还没有实现，大神要是现实了，分享一下哈。
（4）注意训练mobilenet时，在迭代10000次以前，loss和准确率几乎不会提高。一开始我以为是训练代码写错了，后来寻思了很久，才发现是模型太复杂了，所以收敛慢的一比，大概20000次迭代后，准确率才开始蹭蹭的往上长,迭代十万次后准确率才70%。

目录

使用自己的数据集训练MobileNet图像识别（TensorFlow）

1、项目文件结构说明
2、MobileNet的网络：
3、图片数据集
4、制作tfrecords数据格式
5、MobileNet模型
6、训练方法实现过程
7、模型预测
8、其他模型训练方法

1、项目文件结构说明

tensorflow_models_nets：|__dataset   #数据文件    |__record #里面存放record文件    |__train    #train原始图片    |__val      #val原始图片|__models  #保存训练的模型|__slim        #这个是拷贝自slim模块：https://github.com/tensorflow/models/tree/master/research/slim|__test_image #存放测试的图片|__create_labels_files.py #制作trian和val TXT的文件|__create_tf_record.py #制作tfrecord文件|__inception_v1_train_val.py #inception V1的训练文件|__inception_v3_train_val.py # inception V3训练文件|__mobilenet_train_val.py#mobilenet训练文件|__resnet_v1_train_val.py#resnet训练文件|__predict.py # 模型预测文件

2、MobileNet的网络

关于MobileNet模型，请详看这篇博客《轻量级网络--MobileNet论文解读》https://blog.csdn.net/u011974639/article/details/79199306 ，本博客不会纠结于模型原理和论文，主要分享的是用自己的数据集去训练MobileNet的方法。

3、图片数据集

下面是我下载的数据集，共有五类图片，分别是：flower、guitar、animal、houses和plane，每组数据集大概有800张左右。为了照顾网友，下面的数据集，都已经放在Github项目的文件夹dataset上了，不需要你下载了，记得给个“star”哈~

animal：http://www.robots.ox.ac.uk/~vgg/data/pets/ flower：http://www.robots.ox.ac.uk/~vgg/data/flowers/ plane：http://www.robots.ox.ac.uk/~vgg/data/airplanes_side/airplanes_side.tar house：http://www.robots.ox.ac.uk/~vgg/data/houses/houses.tar guitar：http://www.robots.ox.ac.uk/~vgg/data/guitars/guitars.tar 

下载图片数据集后，需要划分为train和val数据集，前者用于训练模型的数据，后者主要用于验证模型。这里提供一个create_labels_files.py脚本，可以直接生成训练train和验证val的数据集txt文件。

#-*-coding:utf-8-*-"""    @Project: googlenet_classification    @File   : create_labels_files.py    @Author : panjq    @E-mail : pan_jinquan@163.com    @Date   : 2018-08-11 10:15:28"""import osimport os.pathdef write_txt(content, filename, mode='w'):    """保存txt数据    :param content:需要保存的数据,type->list    :param filename:文件名    :param mode:读写模式:'w' or 'a'    :return: void    """    with open(filename, mode) as f:        for line in content:            str_line = ""            for col, data in enumerate(line):                if not col == len(line) - 1:                    # 以空格作为分隔符                    str_line = str_line + str(data) + " "                else:                    # 每行最后一个数据用换行符“\n”                    str_line = str_line + str(data) + "\n"            f.write(str_line)def get_files_list(dir):    '''    实现遍历dir目录下,所有文件(包含子文件夹的文件)    :param dir:指定文件夹目录    :return:包含所有文件的列表->list    '''    # parent:父目录, filenames:该目录下所有文件夹,filenames:该目录下的文件名    files_list = []    for parent, dirnames, filenames in os.walk(dir):        for filename in filenames:            # print("parent is: " + parent)            # print("filename is: " + filename)            # print(os.path.join(parent, filename))  # 输出rootdir路径下所有文件（包含子文件）信息            curr_file=parent.split(os.sep)[-1]            if curr_file=='flower':                labels=0            elif curr_file=='guitar':                labels=1            elif curr_file=='animal':                labels=2            elif curr_file=='houses':                labels=3            elif curr_file=='plane':                labels=4            files_list.append([os.path.join(curr_file, filename),labels])    return files_listif __name__ == '__main__':    train_dir = 'dataset/train'    train_txt='dataset/train.txt'    train_data = get_files_list(train_dir)    write_txt(train_data,train_txt,mode='w')    val_dir = 'dataset/val'    val_txt='dataset/val.txt'    val_data = get_files_list(val_dir)    write_txt(val_data,val_txt,mode='w')

注意，上面Python代码，已经定义每组图片对应的标签labels：

flower   ->labels=0guitar   ->labels=1animal  ->labels=2houses ->labels=3plane    ->labels=4

4、制作tfrecords数据格式

有了 train.txt和val.txt数据集，我们就可以制作train.tfrecords和val.tfrecords文件了，项目提供一个用于制作tfrecords数据格式的Python文件：create_tf_record.py，鄙人已经把代码放在另一篇博客：《Tensorflow生成自己的图片数据集TFrecords》https://blog.csdn.net/guyuealian/article/details/80857228 ，代码有详细注释了，所以这里不贴出来了.

注意：
（1）create_tf_record.py将train和val数据分别保存为单个record文件，当图片数据很多时候，会导致单个record文件超级巨大的情况，解决方法就是，将数据分成多个record文件保存，读取时，只需要将多个record文件的路径列表交给“tf.train.string_input_producer”即可。
（2）如何将数据保存为多个record文件呢？请参考鄙人的博客：《Tensorflow生成自己的图片数据集TFrecords》https://blog.csdn.net/guyuealian/article/details/80857228

为了方便大家，项目以及适配了“create_tf_record.py”文件，dataset已经包含了训练和测试的图片，请直接运行create_tf_record.py即可生成tfrecords文件。

对于InceptionNet V1:设置resize_height和resize_width = 224
对于InceptionNet V3:设置resize_height和resize_width = 299

其他模型，请根据输入需要设置resize_height和resize_width的大小

if __name__ == '__main__':    # 参数设置    resize_height = 224  # 指定存储图片高度    resize_width = 224  # 指定存储图片宽度    shuffle=True    log=5    # 产生train.record文件    image_dir='dataset/train'    train_labels = 'dataset/train.txt'  # 图片路径    train_record_output = 'dataset/record/train{}.tfrecords'.format(resize_height)    create_records(image_dir,train_labels, train_record_output, resize_height, resize_width,shuffle,log)    train_nums=get_example_nums(train_record_output)    print("save train example nums={}".format(train_nums))    # 产生val.record文件    image_dir='dataset/val'    val_labels = 'dataset/val.txt'  # 图片路径    val_record_output = 'dataset/record/val{}.tfrecords'.format(resize_height)    create_records(image_dir,val_labels, val_record_output, resize_height, resize_width,shuffle,log)    val_nums=get_example_nums(val_record_output)    print("save val example nums={}".format(val_nums))    # 测试显示函数    # disp_records(train_record_output,resize_height, resize_width)    batch_test(train_record_output,resize_height, resize_width)  create_tf_record.py提供几个重要的函数：create_records()：用于制作records数据的函数，read_records()：用于读取records数据的函数,get_batch_images():用于生成批训练数据的函数get_example_nums：统计tf_records图像的个数(example个数)disp_records(): 解析record文件，并显示图片，主要用于验证生成record文件是否成功

5、MobileNet模型

官网TensorFlow已经提供了使用TF-slim实现的MobileNet模型。

（1）官网模型地址：https://github.com/tensorflow/models/tree/master/research/slim/nets

（2）slim/nets下的模型都是用TF-slim实现的网络结构，关系TF-slim的用法，可参考：
《tensorflow中slim模块api介绍》：https://blog.csdn.net/guvcolie/article/details/77686555

6、训练方法实现过程

训练文件源码已经给了较为详细的注释，不明白请在评论区留言吧

#coding=utf-8import tensorflow as tf import numpy as np import pdbimport osfrom datetime import datetimeimport slim.nets.mobilenet_v1 as mobilenet_v1from create_tf_record import *import tensorflow.contrib.slim as slim'''参考资料：https://www.cnblogs.com/adong7639/p/7942384.html'''labels_nums = 5  # 类别个数batch_size = 16  #resize_height = 224  # mobilenet_v1.default_image_size 指定存储图片高度resize_width = 224   # mobilenet_v1.default_image_size 指定存储图片宽度depths = 3data_shape = [batch_size, resize_height, resize_width, depths]# 定义input_images为图片数据input_images = tf.placeholder(dtype=tf.float32, shape=[None, resize_height, resize_width, depths], name='input')# 定义input_labels为labels数据# input_labels = tf.placeholder(dtype=tf.int32, shape=[None], name='label')input_labels = tf.placeholder(dtype=tf.int32, shape=[None, labels_nums], name='label')# 定义dropout的概率keep_prob = tf.placeholder(tf.float32,name='keep_prob')is_training = tf.placeholder(tf.bool, name='is_training')def net_evaluation(sess,loss,accuracy,val_images_batch,val_labels_batch,val_nums):    val_max_steps = int(val_nums / batch_size)    val_losses = []    val_accs = []    for _ in range(val_max_steps):        val_x, val_y = sess.run([val_images_batch, val_labels_batch])        # print('labels:',val_y)        # val_loss = sess.run(loss, feed_dict={x: val_x, y: val_y, keep_prob: 1.0})        # val_acc = sess.run(accuracy,feed_dict={x: val_x, y: val_y, keep_prob: 1.0})        val_loss,val_acc = sess.run([loss,accuracy], feed_dict={input_images: val_x, input_labels: val_y, keep_prob:1.0, is_training: False})        val_losses.append(val_loss)        val_accs.append(val_acc)    mean_loss = np.array(val_losses, dtype=np.float32).mean()    mean_acc = np.array(val_accs, dtype=np.float32).mean()    return mean_loss, mean_accdef step_train(train_op,loss,accuracy,               train_images_batch,train_labels_batch,train_nums,train_log_step,               val_images_batch,val_labels_batch,val_nums,val_log_step,               snapshot_prefix,snapshot):    '''    循环迭代训练过程    :param train_op: 训练op    :param loss:     loss函数    :param accuracy: 准确率函数    :param train_images_batch: 训练images数据    :param train_labels_batch: 训练labels数据    :param train_nums:         总训练数据    :param train_log_step:   训练log显示间隔    :param val_images_batch: 验证images数据    :param val_labels_batch: 验证labels数据    :param val_nums:         总验证数据    :param val_log_step:     验证log显示间隔    :param snapshot_prefix: 模型保存的路径    :param snapshot:        模型保存间隔    :return: None    '''    saver = tf.train.Saver(max_to_keep=5)    max_acc = 0.0    with tf.Session() as sess:        sess.run(tf.global_variables_initializer())        sess.run(tf.local_variables_initializer())        coord = tf.train.Coordinator()        threads = tf.train.start_queue_runners(sess=sess, coord=coord)        for i in range(max_steps + 1):            batch_input_images, batch_input_labels = sess.run([train_images_batch, train_labels_batch])            _, train_loss = sess.run([train_op, loss], feed_dict={input_images: batch_input_images,                                                                  input_labels: batch_input_labels,                                                                  keep_prob: 0.8, is_training: True})            # train测试(这里仅测试训练集的一个batch)            if i % train_log_step == 0:                train_acc = sess.run(accuracy, feed_dict={input_images: batch_input_images,                                                          input_labels: batch_input_labels,                                                          keep_prob: 1.0, is_training: False})                print("%s: Step [%d]  train Loss : %f, training accuracy :  %g" % (                datetime.now(), i, train_loss, train_acc))            # val测试(测试全部val数据)            if i % val_log_step == 0:                mean_loss, mean_acc = net_evaluation(sess, loss, accuracy, val_images_batch, val_labels_batch, val_nums)                print("%s: Step [%d]  val Loss : %f, val accuracy :  %g" % (datetime.now(), i, mean_loss, mean_acc))            # 模型保存:每迭代snapshot次或者最后一次保存模型            if (i % snapshot == 0 and i > 0) or i == max_steps:                print('-----save:{}-{}'.format(snapshot_prefix, i))                saver.save(sess, snapshot_prefix, global_step=i)            # 保存val准确率最高的模型            if mean_acc > max_acc and mean_acc > 0.7:                max_acc = mean_acc                path = os.path.dirname(snapshot_prefix)                best_models = os.path.join(path, 'best_models_{}_{:.4f}.ckpt'.format(i, max_acc))                print('------save:{}'.format(best_models))                saver.save(sess, best_models)        coord.request_stop()        coord.join(threads)def train(train_record_file,          train_log_step,          train_param,          val_record_file,          val_log_step,          labels_nums,          data_shape,          snapshot,          snapshot_prefix):    '''    :param train_record_file: 训练的tfrecord文件    :param train_log_step: 显示训练过程log信息间隔    :param train_param: train参数    :param val_record_file: 验证的tfrecord文件    :param val_log_step: 显示验证过程log信息间隔    :param val_param: val参数    :param labels_nums: labels数    :param data_shape: 输入数据shape    :param snapshot: 保存模型间隔    :param snapshot_prefix: 保存模型文件的前缀名    :return:    '''    [base_lr,max_steps]=train_param    [batch_size,resize_height,resize_width,depths]=data_shape    # 获得训练和测试的样本数    train_nums=get_example_nums(train_record_file)    val_nums=get_example_nums(val_record_file)    print('train nums:%d,val nums:%d'%(train_nums,val_nums))    # 从record中读取图片和labels数据    # train数据,训练数据一般要求打乱顺序shuffle=True    train_images, train_labels = read_records(train_record_file, resize_height, resize_width, type='normalization')    train_images_batch, train_labels_batch = get_batch_images(train_images, train_labels,                                                              batch_size=batch_size, labels_nums=labels_nums,                                                              one_hot=True, shuffle=True)    # val数据,验证数据可以不需要打乱数据    val_images, val_labels = read_records(val_record_file, resize_height, resize_width, type='normalization')    val_images_batch, val_labels_batch = get_batch_images(val_images, val_labels,                                                          batch_size=batch_size, labels_nums=labels_nums,                                                          one_hot=True, shuffle=False)    # Define the model:    with slim.arg_scope(mobilenet_v1.mobilenet_v1_arg_scope()):        out, end_points = mobilenet_v1.mobilenet_v1(inputs=input_images, num_classes=labels_nums,                                                    dropout_keep_prob=keep_prob, is_training=is_training,                                                    global_pool=True)        # Specify the loss function: tf.losses定义的loss函数都会自动添加到loss函数,不需要add_loss()了    tf.losses.softmax_cross_entropy(onehot_labels=input_labels, logits=out)  # 添加交叉熵损失loss=1.6    # slim.losses.add_loss(my_loss)    loss = tf.losses.get_total_loss(add_regularization_losses=True)  # 添加正则化损失loss=2.2    # Specify the optimization scheme:    # 在定义训练的时候, 注意到我们使用了`batch_norm`层时,需要更新每一层的`average`和`variance`参数,    # 更新的过程不包含在正常的训练过程中, 需要我们去手动像下面这样更新    # 通过`tf.get_collection`获得所有需要更新的`op`    update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)    # 使用`tensorflow`的控制流, 先执行更新算子, 再执行训练    with tf.control_dependencies(update_ops):        print("update_ops:{}".format(update_ops))        # create_train_op that ensures that when we evaluate it to get the loss,        # the update_ops are done and the gradient updates are computed.        # train_op = tf.train.MomentumOptimizer(learning_rate=base_lr, momentum=0.9).minimize(loss)        train_op = tf.train.AdadeltaOptimizer(learning_rate=base_lr).minimize(loss)    accuracy = tf.reduce_mean(tf.cast(tf.equal(tf.argmax(out, 1), tf.argmax(input_labels, 1)), tf.float32))    # 循环迭代过程    step_train(train_op=train_op, loss=loss, accuracy=accuracy,               train_images_batch=train_images_batch,               train_labels_batch=train_labels_batch,               train_nums=train_nums,               train_log_step=train_log_step,               val_images_batch=val_images_batch,               val_labels_batch=val_labels_batch,               val_nums=val_nums,               val_log_step=val_log_step,               snapshot_prefix=snapshot_prefix,               snapshot=snapshot)if __name__ == '__main__':    train_record_file='dataset/record/train224.tfrecords'    val_record_file='dataset/record/val224.tfrecords'    train_log_step=100    base_lr = 0.001  # 学习率    # 重头开始训练的话，mobilenet收敛慢的一比，大概20000次迭代后，准确率开始蹭蹭的往上长,迭代十万次后准确率才70%    max_steps = 100000  # 迭代次数    train_param=[base_lr,max_steps]    val_log_step=500    snapshot=2000#保存文件间隔    snapshot_prefix='models/model.ckpt'    train(train_record_file=train_record_file,          train_log_step=train_log_step,          train_param=train_param,          val_record_file=val_record_file,          val_log_step=val_log_step,          labels_nums=labels_nums,          data_shape=data_shape,          snapshot=snapshot,          snapshot_prefix=snapshot_prefix)

7、模型预测

模型预测，项目只提供一个predict.py，实质上，你只需要稍微改改，就可以预测其他模型

#coding=utf-8import tensorflow as tf import numpy as np import pdbimport cv2import osimport globimport slim.nets.inception_v3 as inception_v3from create_tf_record import *import tensorflow.contrib.slim as slimdef  predict(models_path,image_dir,labels_filename,labels_nums, data_format):    [batch_size, resize_height, resize_width, depths] = data_format    labels = np.loadtxt(labels_filename, str, delimiter='\t')    input_images = tf.placeholder(dtype=tf.float32, shape=[None, resize_height, resize_width, depths], name='input')    #其他模型预测请修改这里    with slim.arg_scope(inception_v3.inception_v3_arg_scope()):        out, end_points = inception_v3.inception_v3(inputs=input_images, num_classes=labels_nums, dropout_keep_prob=1.0, is_training=False)    # 将输出结果进行softmax分布,再求最大概率所属类别    score = tf.nn.softmax(out,name='pre')    class_id = tf.argmax(score, 1)    sess = tf.InteractiveSession()    sess.run(tf.global_variables_initializer())    saver = tf.train.Saver()    saver.restore(sess, models_path)    images_list=glob.glob(os.path.join(image_dir,'*.jpg'))    for image_path in images_list:        im=read_image(image_path,resize_height,resize_width,normalization=True)        im=im[np.newaxis,:]        #pred = sess.run(f_cls, feed_dict={x:im, keep_prob:1.0})        pre_score,pre_label = sess.run([score,class_id], feed_dict={input_images:im})        max_score=pre_score[0,pre_label]        print("{} is: pre labels:{},name:{} score: {}".format(image_path,pre_label,labels[pre_label], max_score))    sess.close()if __name__ == '__main__':    class_nums=5    image_dir='test_image'    labels_filename='dataset/label.txt'    models_path='models/model.ckpt-10000'    batch_size = 1  #    resize_height = 299  # 指定存储图片高度    resize_width = 299  # 指定存储图片宽度    depths=3    data_format=[batch_size,resize_height,resize_width,depths]    predict(models_path,image_dir, labels_filename, class_nums, data_format)

8、其他模型训练方法

上面的程序是训练MobileNet的完整过程，实质上，稍微改改就可以支持训练 inception V1，V2和resnet 啦，改动方法也很简单，以 MobileNe训练代码改为resnet_v1模型为例：

（1）import 改为：

# 将import slim.nets.mobilenet_v1 as mobilenet_v1# 改为import slim.nets.resnet_v1 as resnet_v1

（2）record数据

制作record数据时，需要根据模型输入设置：

resize_height = 224  # 指定存储图片高度resize_width = 224  # 指定存储图片宽度

（3）定义模型和默认参数修改：

    # 将    # Define the model:    with slim.arg_scope(mobilenet_v1.mobilenet_v1_arg_scope()):        out, end_points = mobilenet_v1.mobilenet_v1(inputs=input_images, num_classes=labels_nums,                                                    dropout_keep_prob=keep_prob, is_training=is_training,                                                    global_pool=True)    # 改为    # Define the model:    with slim.arg_scope(resnet_v1.resnet_arg_scope()):        out, end_points = resnet_v1.resnet_v1_101(inputs=input_images, num_classes=labels_nums, is_training=is_training,global_pool=True)

（4）修改优化方案

对于大型的网络模型，重头开始训练，是很难收敛的。训练mobilenet时，在迭代10000次以前，loss和准确率几乎不会提高。一开始我以为是训练代码写错了，后来寻思了很久，才发现是模型太复杂了，所以收敛慢的一比，大概20000次迭代后，准确率才开始蹭蹭的往上长,迭代十万次后准确率才70%，若训练过程发现不收敛，请尝试修改：

1、等！！！！至少你要迭代50000次，才能说你的模型不收敛！
2、增大或减小学习率参数：base_lr（个人经验：模型越深越复杂时，学习率越小）
3、改变优化方案：如使用MomentumOptimizer或者AdadeltaOptimizer等优化方法
4、是否有设置默认的模型参数：如slim.arg_scope(inception_v1.inception_v1_arg_scope())

……最后，就可以Train了！是的，就是那么简单~

技术的道路一个人走着极为艰难？

一身的本领得不施展？

优质的文章得不到曝光？

别担心，

即刻起，CSDN 将为你带来创新创造创变展现的大舞台，

扫描下方二维码，欢迎加入 CSDN 「原力计划」！

推荐阅读

• 大四学生发明文言文编程语言，设计思路清奇

• 超模脸、网红脸、萌娃脸...换头像不重样？我开源了5款人脸生成器

• 解读 | 2019年10篇计算机视觉精选论文（上）

• 高通：2 亿像素手机 2020 年诞生！

• 英特尔首推异构编程神器 oneAPI，可让程序员少加班！

• VS Code 成主宰、Vue 备受热捧！2019 前端开发趋势必读

• 我在华为做外包的真实经历

• 2019 区块链大事记 | Libra 横空出世，莱特币减产，美国放行 Bakkt……这一年太精彩！

• 互联网诞生记: 浪成于微澜之间

• 

  你点的每个“在看”，我都认真当成了AI