用tensorflow训练自己的图片集-用TFRecords将代码导入神经网络

来源:互联网 发布:衣阿华 知乎 编辑:程序博客网 时间:2024/06/06 04:03

           之前用过Alexnet,vggnet都成功地跑了cifar-10的数据集,现在想要训练自己的数据。第一个难点就是mnist和cifar10都是已经打好标签压缩好的图片集,用一条命令就能将img和label读出并使用。而训练自己的图像需要自己打上标签并传入神经网络。被困在这里两天多,今天终于搞懂了,其实使用TFRecords很简单。并且不用自己制作标签,非常的方便。

1.制作自己的训练集TFRecords

       首先需要把你的图片根据分类放在不同的文件夹下,class的名字就是你要分类的名字。我是要识别fsk波形和qpsk波形,就分别命名为fsk和qpsk。像这样

       

       

       接下来就可以写代码了

        

import osimport tensorflow as tffrom PIL import Imageimport matplotlib.pyplot as pltimport numpy as npcwd = '/home/modric/Downloads/test/'classes = {'fsk','qpsk'}writer = tf.python_io.TFRecordWriter("train.tfrecords")for index, name in enumerate(classes):    class_path = cwd + name + '/'    for img_name in os.listdir(class_path):        img_path = class_path + img_name        img = Image.open(img_path)        img = img.resize((128,128))        img_raw = img.tobytes()        example = tf.train.Example(features = tf.train.Features(feature={            "label":tf.train.Feature(int64_list = tf.train.Int64List(value=[index])),            'img_raw':tf.train.Feature(bytes_list = tf.train.BytesList(value=[img_raw]))        }))        writer.write(example.SerializeToString())writer.close()

2.使用队列读取img和label

       一旦生成了TFRecords文件,为了高效地读取数据,TF中使用队列(queue)读取数据。

def read_and_decode(filename):#read .tfrecords    filename_queue = tf.train.string_input_producer([filename])#create a queue    reader = tf.TFRecordeReader()    _, serialized_example = reader.read(filename_queue)#return filename and file    features = tf.parse_single_example(serialized_example, features = {                                                               'label':tf.FixedLenFeature([], tf.int64),                                                               'img_raw':tf.FixedLenFeature([], tf.string),                                                           })#take out image and label    img = tf.decode_raw(features['img_raw']. tf.uint8)    img = tf.reshape(img, [128, 128, 3])#reshape to 128,128,3    img = tf.cast(img, tf.float32)*(1./255) - 0.5#throw out the img tensor    label = tf.cast(features['label'], tf.int32)    return img, label

我之前就是卡在class那了,不知道应该怎么操作,文件也不知道命名的规则和存放在哪里。后来上手实践过后发现还是比较简单,所以还是要动手练呀,生命在于运动,加油!

多谢这位仁兄的文章http://blog.csdn.net/u012759136/article/details/52232266

阅读全文
1 0