tfrecords 格式数据训练mnist

import os
import tensorflow as tf
from PIL import Image
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
def data_to_tfrecord(images, labels, filename):
# Save data into TFRecord
    if os.path.isfile(filename):
        print("%s exists" % filename)
        return
    print("Converting data into %s ..." % filename)
    cwd = os.getcwd()
    writer = tf.python_io.TFRecordWriter(filename)
    for index, img_name in enumerate(images):
        print(index)
        img = Image.open(img_name)
        img = img.resize((28, 28))
        img_raw = img.tobytes()
 
        example = tf.train.Example(features=tf.train.Features(feature={
"label": tf.train.Feature(int64_list=tf.train.Int64List(value=[int(labels[index])])),
            'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),
        }))
        writer.write(example.SerializeToString())  # Serialize To String
    writer.close()
 
def tfrecord_to_data(filename):
    # generate a queue with a given file name
    print("reading tfrecords from {}".format(filename))
    filename_queue = tf.train.string_input_producer([filename])
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)  
    features = tf.parse_single_example(serialized_example,features={
'label': tf.FixedLenFeature([], tf.int64),
'img_raw': tf.FixedLenFeature([], tf.string),
})
    img = tf.decode_raw(features['img_raw'], tf.uint8)
    img = tf.reshape(img, [28, 28, 3])
    img = tf.cast(img, tf.float32) * (1. / 255) - 0.5
    label = tf.cast(features['label'], tf.int64)
    return img, label
 
 
 
def decode_tfrecords(filename):
#show the tfrecords
    for serialized_example in tf.python_io.tf_record_iterator(filename):
        example = tf.train.Example()
        example.ParseFromString(serialized_example)
 
        image = example.features.feature['img_raw'].bytes_list.value
        label = example.features.feature['label'].int64_list.value
        print(image, label)
 
 
def read_data_from_paths(file_path, name):
    labels = []
    file_names = []
 
    file_name = os.path.join(file_path, name)
    train_txt = open(file_name,'r')
 
    for idx in train_txt:
        idx=idx.rstrip('\n')
        spt = idx.split(' ')
        file_names.append(os.path.join(file_path, spt[0]))
        labels.append(spt[1])
    return file_names, labels
 
 
def train():
    #network
    batch_size = 64
    inputs = tf.placeholder(tf.float32, [batch_size, 28, 28, 3], name='inputs')
    conv1 = tf.layers.conv2d(inputs=inputs, filters=64, kernel_size=(3, 3), padding="same", activation=None)
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
    conv2 = tf.layers.conv2d(inputs=pool1, filters=128, kernel_size=(3, 3), padding="same", activation=None)
    pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
 
    pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 128])
    fc1 = tf.layers.dense(pool2_flat, 500, activation=tf.nn.relu)
    fc2 = tf.layers.dense(fc1, 10, activation=tf.nn.relu)
    y_out = tf.nn.softmax(fc2)
 
    y_ = tf.placeholder(tf.float32, [batch_size, 10])
    cross_entropy = -tf.reduce_mean(y_ * tf.log(y_out))  # 计算交叉熵
 
    learning_rate=1e-3
    train_step = tf.train.AdamOptimizer(learning_rate).minimize(cross_entropy)
    correct_prediction = tf.equal(tf.argmax(y_out, 1), tf.argmax(y_, 1))  
accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))
 
    img, label = tfrecord_to_data("./mnist_train.tfrecords")
    img_batch, label_batch = tf.train.shuffle_batch([img, label],batch_size=batch_size, capacity=2000,min_after_dequeue=1000)
 
    init = tf.global_variables_initializer()
 
    with tf.Session() as session:
        session.run(init)
        threads = tf.train.start_queue_runners()
        saver = tf.train.Saver(tf.global_variables(), max_to_keep=20)
        for i in range(400):
            img_batch_i, label_batch_i = session.run([img_batch, tf.one_hot(label_batch, depth=10)])
 
            feed = {inputs: img_batch_i, y_: label_batch_i}
            loss,_,acc=session.run([cross_entropy,train_step,accuracy], feed_dict=feed)
 
            print("step%d loss:%f accuracy:%F"%(i,loss,acc))
            if i>100:
                learning_rate=learning_rate*0.1
        saver.save(session, "./save/mnist.ckpt")
 
 
if __name__=="__main__":
    #function1,image to tfrecords,not use comment
    file_path=""
    name="test.txt"
    file_names, labels=read_data_from_paths(file_path, name)
    tfrecord_name="mnist_test.tfrecords"
    data_to_tfrecord(file_names, labels, tfrecord_name)
 
    #function2,decode tfrecords,not use comment
    decode_tfrecords("./mnist_test.tfrecords")
    decode_tfrecords("E:/ocr_pdf_rec/crnn_v2.2/train/data/test_2k.tfrecords")
 
    #function3,tfrecords to images,not use comment
    train()