迁移学习用于图像识别的Tensorflow实现

最近在研究目标识别的YOLO论文，想按照论文中的模型进行实现，不过发现按照论文中的24层CNN网络结构，在我的GTX750Ti 2G显存的卡上没法跑起来，看来是时候要换张大容量的显卡了。不过在换显卡之前，我想先测试一下在现有的显卡基础上，是否有其他办法可以跑YOLO算法。其中一个办法是采用迁移学习的思路，把别人训练好的卷积网络直接拿来计算图像的特征。这样我就不用自己搭建这么多层网络来重新训练了。

在Tensorflow Hub里面有很多训练好的卷积神经网络，这些网络大部分都是以Imagenet 2012竞赛的1000类图像来进行训练的，这个数据集大概有146GB，按照YOLO论文作者的说法，他的模型也是先在这个数据集上跑了1个星期来进行预训练的，可想而知如果自己用这个数据集来训练的话，需要有多大的显卡资源才能满足要求。因此明智的做法是直接采用Tensorflow Hub里面的模型来进行迁移学习。我选取了其中的Google Inception V3模型，这个也是一个非常出名的模型，有着很高的准确度。在做迁移学习到YOLO模型的搭建之前，我先用这个模型来进行图像的重新训练识别，先练一下手，熟悉一下迁移学习的做法。

我也是采用了Tensorflow网站的Image Retrain Guide里面的Flowers数据集，但是没有按照Guide里面的程序来做，而是自己写了一个程序来实现，这也是更好的学习掌握Tensorflow的一个好方法。

我的程序分为两个部分。第一部分是把Flowers数据集转换为TFRECORD格式，每个图片都按照Inception V3模型的要求缩放为299×299像素，并把类型标签也一并写入到文件中。这样方便我们进行训练时读取数据。这个程序代码如下所示：

import os
import cv2
import tensorflow as tf
import numpy as np

reshape_width = 299.0
reshape_height = 299.0

def make_example(image, label):
    return tf.train.Example(features=tf.train.Features(feature={
        'image' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),
        'label' : tf.train.Feature(bytes_list=tf.train.BytesList(value=[label]))
    }))

flower_classes = {"daisy":0, "dandelion":1, "roses":2, "sunflowers":3, "tulips":4}

for flower_class in flower_classes.keys():
    writer = tf.python_io.TFRecordWriter(flower_class+".tfrecord")
    folder_path = "/home/roy/flower_photos/"+flower_class
    files = os.listdir(folder_path)
    label = np.array([flower_classes[flower_class]])
    for jpgfile in files:
        img = cv2.imread(folder_path+"/"+jpgfile, cv2.IMREAD_COLOR)
        img = cv2.resize(img, (int(reshape_width), int(reshape_height)))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.uint8)
        ex = make_example(img.tobytes(), label.tobytes())
        writer.write(ex.SerializeToString())
    writer.close()

程序运行完后，会生成5个TFRECORD文件，每个文件对应一种花的类型。

未完待续。。。

迁移学习用于图像识别的Tensorflow实现

猜你喜欢