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


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

1.制作自己的训练集TFRecords

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

       

       

       接下来就可以写代码了

        

import os
import tensorflow as tf
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

cwd = '/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

你可能感兴趣的:(TensorFlow)