使用tf-record制作自己的数据集用于训练

1 准备数据集，准备了5个类别的数据，写了一个小程序，写入到 对应文件夹下：

import os
import shutil
path = r"C:\Users\HUANG\Desktop\DATA\train"
test_path = r"C:\Users\HUANG\Desktop\DATA\test"
car = r"C:\Users\HUANG\Desktop\DATA\car"
dinosaur = r"C:\Users\HUANG\Desktop\DATA/dinosaur"
elephant = r"C:\Users\HUANG\Desktop\DATA\dinosaur"
flower = r"C:\Users\HUANG\Desktop\DATA\flower"
horse = r"C:\Users\HUANG\Desktop\DATA\horse"
classes ={'car', 'dinosaur', 'elephant','flower', 'horse'}

#files = os.listdir(path)
#
##print(files)
#for file in files:
##    print(file.)
#    #print(type(file))
#    string = int(file.split('.')[0])
#    #print(string)
#    for index , name in enumerate(classes):
#        i = index+3
#        print(i, name)
#        if string // 100 == i:
#        #print(1)
#        #car = open(r"C:\Users\HUANG\Desktop\DATA\car", 'w')
#        #car.write(file)
#            shutil.move(os.path.join(path, file), name)
#    print('完成')


file1 = os.listdir(test_path)

for file in file1:
    string = int(file.split('.')[0])
    if string // 100 ==7:
        shutil.move(os.path.join(test_path, file), horse)
    print('完成 ！')
    

2 使用tf-record制作自己的数据集，以及数据的读取，代码如下：

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


path = r"C:\Users\HUANG\Desktop\DATA"
train_record_path = r"C:\Users\HUANG\Desktop\DATA/train.tfrecords"
test_record_path = r"C:\Users\HUANG\Desktop\DATA/test.tfrecords"
classes = {'car', 'dinosaur', 'elephant', 'flower', 'horse'}

def _byteslist(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _int64list(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def creat_train_record():
    writer = tf.python_io.TFRecordWriter(train_record_path)
    NUM = 1
    for index, name in enumerate(classes):
        class_path = path + "/" + name + '/'
        l = int(len(os.listdir(class_path))*0.7)
        for image_name in os.listdir(class_path)[:l]:
            image_path = class_path + image_name
            img = Image.open(image_path)
            img = img.resize((128, 128))
            img_raw = img.tobytes()
            example = tf.train.Example(
                    features=tf.train.Features(feature={
                            'label': _int64list(index),
                            'img_raw': _byteslist(img_raw)}))
            writer.write(example.SerializeToString())
            print('creat train record in ', NUM)
            NUM += 1
    writer.close()
    print('creat_train_record success !')
    
def creat_test_record():
    writer = tf.python_io.TFRecordWriter(test_record_path)
    NUM = 1
    for index, name in enumerate(classes):
        class_path = path + '/' + name +'/'
        l = int(len(os.listdir(class_path))*0.7)
        for image_name in os.listdir(class_path)[l:]:
            image_path = class_path + image_name
            img = Image.open(image_path)
            img = img.resize((128,128))
            img_raw = img.tobytes()
            example = tf.train.Example(
                    features=tf.train.Features(feature={
                                               'label': _int64list(index),
                                               'img_raw': _byteslist(img_raw)}))
            writer.write(example.SerializeToString())
            print('creat test record in', NUM)
            NUM += 1
    writer.close()
    print('creat test record finished !')

def read_record(filename, img_w, img_h):
    filename_queue = tf.train.string_input_producer([filename])
    reader = tf.TFRecordReader()
    _, serialize_example = reader.read(filename_queue)
    feature = tf.parse_single_example(
            serialize_example,
            features={
                    'label': tf.FixedLenFeature([], tf.int64),
                    'img_raw': tf.FixedLenFeature([], tf.string)})
    label = feature['label']
    img = feature['img_raw']
    img = tf.decode_raw(img, tf.uint8)
    img = tf.reshape(img, (128, 128, 3))
    img = tf.image.resize_image_with_crop_or_pad(img, img_w, img_h) #裁剪图片
    img = tf.cast(img, tf.float32)/255
    label = tf.cast(label, tf.int32)
    return img, label

def get_batch_record(filename, batch_size, img_W, img_H):
    image, label = read_record(filename, img_W, img_H)
    image_batch, label_batch= tf.train.shuffle_batch([image, label],
                                                     batch_size=batch_size,
                                                     capacity=200+batch_size*3,
                                                     min_after_dequeue=1)
    label_batch = tf.one_hot(label_batch,depth=5)
    return image_batch, label_batch

#img, label = get_batch_record(train_record_path, 3, 256, 256)
'''
with tf.Session() as sess:
    #op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer())
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess, coord)
    for i in range(20):
        image, l =sess.run([img, label])
        print(image[0].shape)
        plt.imshow(image[0])
        plt.show()
    coord.request_stop()
'''   




'''    
    for i in range(2):
        try:
            while not coord.should_stop():
                im, l = sess.run([img, label])
                print(im.shape)
        except tf.errors.OutOfRangeError:
            print('error')
        finally:
            coord.request_stop()
    coord.join(threads)

try:
    while not coord.should_stop():
        # Run training steps or whatever
        sess.run(train_op)
 
except tf.errors.OutOfRangeError:
    print 'Done training -- epoch limit reached'
finally:
    # When done, ask the threads to stop.
    coord.request_stop()
 
# Wait for threads to finish.
coord.join(threads)
'''        
       
    

3：搭建网络，代码如下：

'''
model 输入为24x128x128x3
'''
import tensorflow as tf

def weight_variable(shape):
    init = tf.truncated_normal(shape=shape, stddev =0.1, dtype= tf.float32)
    return tf.Variable(init)

def bias_variable(shape):
    init = tf.constant(0.1, shape = shape)
    return tf.Variable(init)

def conv2d(x, w):
    conv = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding ='SAME')
    return conv

def max_pool_2x2(x, name):
    return tf.nn.max_pool(x, strides=[1,2,2,1], ksize=[1,2,2,1], padding='SAME', name=name)

def avg_pool_8x8(x, name):
    return tf.nn.avg_pool(x, strides=[1,8,8,1], ksize=[1,8,8,1], padding='SAME')


def CNN(image, batch_size, num_classes):
    #第一卷积层 （3x3x3x32）
    with tf.variable_scope('conv1') as scope:
        w_conv1 = weight_variable([3,3,3,32])
        b_conv1 = bias_variable([32])
        h_conv1 = tf.nn.relu(conv2d(image, w_conv1)+b_conv1, name='conv1')
        print(h_conv1.shape)  #(?, 128, 128, 32)
    
    #第一池化层
    with tf.variable_scope('pool1') as scope:
        h_pool1 = max_pool_2x2(h_conv1, name ='pool')
        print(h_pool1.shape)  #(?, 64, 64, 32)
    
    #第二卷积层 (3x3x32x64)
    with tf.variable_scope('conv2') as scope:
        w_conv2 = weight_variable([3,3,32,64])
        b_conv2 = bias_variable([64])
        h_conv2 = tf.nn.relu(conv2d(h_pool1, w_conv2) + b_conv2, name='conv2')
        print(h_conv2.shape)  # (?, 64, 64, 64)
    #第二池化层
    with tf.variable_scope('pool2') as scope:
        h_pool2 = max_pool_2x2(h_conv2,name ='pool')
        print(h_pool2.shape)  #(?, 32, 32, 64)
    
    #第三卷积层
    with tf.variable_scope('conv3') as scope:
        w_conv3 = weight_variable([3,3,64,32])
        b_conv3 = bias_variable([32])
        h_conv3 = tf.nn.relu(conv2d(h_pool2, w_conv3)+b_conv3, name='conv3')
        print(h_conv3.shape)  #(?, 32, 32, 32)
    
    #第三池化层
    with tf.variable_scope('pool3') as scope:
        h_pool3 = max_pool_2x2(h_conv3, name='pool')
        print(h_pool3.shape)  #(?, 16, 16, 32)
    
    #第四卷积层
    with tf.variable_scope('conv4') as scope:
        w_conv4 = weight_variable([3,3,32,16])
        b_conv4 = bias_variable([16])
        h_conv4 = tf.nn.relu(conv2d(h_pool3, w_conv4)+b_conv4, name='conv4')
        print(h_conv4.shape)  #(?, 16, 16, 16)
    
    with tf.variable_scope('pool5') as scope:
        h_pool4 = max_pool_2x2(h_conv4, name='pool')
        print(h_pool4.shape)  #(?, 8, 8, 16)
    
     #第五卷积层
    with tf.variable_scope('conv5') as scope:
        w_conv5 = weight_variable([3,3,16,5])
        b_conv5 = bias_variable([5])
        h_conv5 = tf.nn.relu(conv2d(h_pool4, w_conv5)+b_conv5, name='conv4')
        print(h_conv5.shape)   #(?,8,8,5)
    
        
    #第五池化层
    with tf.variable_scope('pool5') as scope:
        h_pool5 = avg_pool_8x8(h_conv5, name='pool')  #batch_sizex1x1x16
        print(h_pool5.shape)  #(?, 1,1,5) 
    
    #全连接层
    
#    with tf.variable_scope('fc') as scope:
#        w_fc5 = weight_variable([batch_size, 1,1, num_classes])
#        b_fc5 = bias_variable([num_classes])
#        h_out = tf.matmul(h_pool4, w_fc5)+b_fc5
#        print(h_out.shape)
    # softmax层
    out_hat = tf.reshape(h_pool5, [-1, num_classes])
    pred = tf.nn.softmax(out_hat)
    print(pred.shape) # (?,5)
    return pred

def loss(pred, label):
    with tf.variable_scope('loss') as scope:
        loss = -tf.reduce_sum(label*tf.log(pred))
    return loss

def optimizer(loss, learning_rate, global_step):
    with tf.variable_scope('optimizer') as scope:
        optimizer = tf.train.AdamOptimizer(learning_rate).minimize(loss, global_step=global_step)
    return optimizer

def accuracy(logits, labels):
    with tf.variable_scope('accuracy') as scope:
        correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(labels, 1))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    return accuracy
        
    
    
        

4：开始训练，代码如下：

'''
train
'''
import tensorflow as tf
import matplotlib.pyplot as plt
from generate1 import get_batch_record
from model import CNN, loss, optimizer, accuracy
import numpy as np

filename = r"C:\Users\HUANG\Desktop\DATA/train.tfrecords"
batch_size =24
num_classes = 5
img_w = 128
img_h = 128
learning_rate = 0.001
train_epoch =2000
display_step = 100

x = tf.placeholder(tf.float32, [None, img_w, img_h, 3])
y = tf.placeholder(tf.float32, [None, num_classes])

global_step = tf.Variable(0,trainable=None)

decay_learning_rate=tf.train.exponential_decay(learning_rate, global_step=global_step, decay_steps=100, decay_rate=0.9)

pred = CNN(x, batch_size, num_classes)
loss = loss(pred, y)
optimizer = optimizer(loss, decay_learning_rate, global_step)
accuracy = accuracy(pred, y)
image_batch, label_batch = get_batch_record(filename, batch_size, img_w, img_h)
fig_loss = []
fig_acc= []

with tf.Session() as sess:
    sess.run(tf.local_variables_initializer())
    sess.run(tf.global_variables_initializer())
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess, coord)
    print('开始训练')
    for i in range(train_epoch):
        image, label = sess.run([image_batch, label_batch])
        _, l, acc = sess.run([optimizer, loss, accuracy], 
                             feed_dict={x:image, y:label})
        if i % display_step ==0:
            fig_loss.append(l)
            fig_acc.append(acc)
            print('epoch:', '%04d'%(i+1), 'loss:' '{:.4f}'.format(l), 'acc:', '%.4f'%(acc))
    
    fig,ax1=plt.subplots()
    ax2=ax1.twinx()
    ln1=plt.plot(np.arange(len(fig_loss)), fig_loss, 'r', label='loss')
    ln2=plt.plot(np.arange(len(fig_acc)), fig_acc, 'g', label='accuracy')
    ax1.set_xlabel('iteration')
    ax2.set_ylabel('training loss')
    ax2.set_ylabel('training accuracy')
    
    lns = ln1 + ln2
    labels = ["Loss", "Accuracy"]
    # labels = [l.get_label() for l in lns]
    plt.legend(lns, labels, loc='upper left')
    plt.grid(True)
    
    
    plt.figure(2)
    plt.plot(np.arange(len(fig_loss)), fig_loss, 'r', label='loss')
    plt.legend(labels=['loss'])
    plt.grid(True)
    
    plt.figure(3)
    plt.plot(np.arange(len(fig_acc)), fig_acc,'g',label='accuracy')
    plt.legend(labels=["acuracy"])
    
    plt.grid(True)


    plt.show()

5：结果展示：

注：由于主要目的是了解tf-reocrd制作数据集以及数据读取的过程，网络结构搭建较为随意，训练准确率基本到达100%，主要是熟悉一个数据制作、读取、训练的过程，因此没有进一步验证测试集的情况，有兴趣的可以自行加上。

以下为训练结果：

开始训练
epoch: 0001 loss:41.0609 acc: 0.0000
epoch: 0101 loss:16.8334 acc: 0.7500
epoch: 0201 loss:5.1557 acc: 0.9583
epoch: 0301 loss:3.0682 acc: 0.9583
epoch: 0401 loss:1.1848 acc: 1.0000
epoch: 0501 loss:5.9951 acc: 0.9167
epoch: 0601 loss:0.4838 acc: 1.0000
epoch: 0701 loss:0.2849 acc: 1.0000
epoch: 0801 loss:0.3520 acc: 1.0000
epoch: 0901 loss:0.4722 acc: 1.0000
epoch: 1001 loss:0.0931 acc: 1.0000
epoch: 1101 loss:0.1719 acc: 1.0000
epoch: 1201 loss:0.0509 acc: 1.0000
epoch: 1301 loss:0.1252 acc: 1.0000
epoch: 1401 loss:0.0354 acc: 1.0000
epoch: 1501 loss:0.1840 acc: 1.0000
epoch: 1601 loss:0.1043 acc: 1.0000
epoch: 1701 loss:0.0568 acc: 1.0000
epoch: 1801 loss:0.1443 acc: 1.0000
epoch: 1901 loss:0.0242 acc: 1.0000

6  文件说明：

数据文件分别为：car，dinosuar, elephant, flower, horse

train.records 为制作的训练文件，

test.record为制作的测试文件

generate1.py ——制作tf-record文件，以及数据的读取

model.py  ——网络结构

train.py——训练