Tensorflow学习教程------读取数据、建立网络、训练模型，小巧而完整的代码示例

紧接上篇Tensorflow学习教程------tfrecords数据格式生成与读取，本篇将数据读取、建立网络以及模型训练整理成一个小样例，完整代码如下。

#coding:utf-8
import tensorflow as tf
import os
def read_and_decode(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, [227, 227, 3])
    img = (tf.cast(img, tf.float32) * (1. / 255) - 0.5)*2
    label = tf.cast(features['label'], tf.int32)
    print img,label
    return img, label

def get_batch(image, label, batch_size,crop_size):
    #数据扩充变换
    distorted_image = tf.random_crop(image, [crop_size, crop_size, 3])#随机裁剪
    distorted_image = tf.image.random_flip_up_down(distorted_image)#上下随机翻转
    distorted_image = tf.image.random_brightness(distorted_image,max_delta=63)#亮度变化
    distorted_image = tf.image.random_contrast(distorted_image,lower=0.2, upper=1.8)#对比度变化  

    #生成batch
    #shuffle_batch的参数：capacity用于定义shuttle的范围，如果是对整个训练数据集，获取batch，那么capacity就应该够大
    #保证数据打的足够乱
    images, label_batch = tf.train.shuffle_batch([distorted_image, label],batch_size=batch_size,
                                                 num_threads=1,capacity=2000,min_after_dequeue=1000) 

    return images, label_batch

class network(object):
    #构造函数初始化 卷积层 全连接层
    def __init__(self):
        with tf.variable_scope("weights"):
           self.weights={  

                'conv1':tf.get_variable('conv1',[4,4,3,20],initializer=tf.contrib.layers.xavier_initializer_conv2d()),  

                'conv2':tf.get_variable('conv2',[3,3,20,40],initializer=tf.contrib.layers.xavier_initializer_conv2d()),  

                'conv3':tf.get_variable('conv3',[3,3,40,60],initializer=tf.contrib.layers.xavier_initializer_conv2d()), 

                'fc1':tf.get_variable('fc1',[3*3*60,120],initializer=tf.contrib.layers.xavier_initializer()),
                'fc2':tf.get_variable('fc2',[120,2],initializer=tf.contrib.layers.xavier_initializer()),  

                }
        with tf.variable_scope("biases"):
            self.biases={
                'conv1':tf.get_variable('conv1',[20,],initializer=tf.constant_initializer(value=0.0, dtype=tf.float32)),
                'conv2':tf.get_variable('conv2',[40,],initializer=tf.constant_initializer(value=0.0, dtype=tf.float32)),
                'conv3':tf.get_variable('conv3',[60,],initializer=tf.constant_initializer(value=0.0, dtype=tf.float32)),  

                'fc1':tf.get_variable('fc1',[120,],initializer=tf.constant_initializer(value=0.0, dtype=tf.float32)),
                'fc2':tf.get_variable('fc2',[2,],initializer=tf.constant_initializer(value=0.0, dtype=tf.float32)), 

            }

    def buildnet(self,images):
        #向量转为矩阵
        images = tf.reshape(images, shape=[-1, 39,39, 3])# [batch, in_height, in_width, in_channels]
        images=(tf.cast(images,tf.float32)/255.-0.5)*2#归一化处理  

        #第一层
        conv1=tf.nn.bias_add(tf.nn.conv2d(images, self.weights['conv1'], strides=[1, 1, 1, 1], padding='SAME'),
                             self.biases['conv1'])
        relu1= tf.nn.relu(conv1)
        pool1=tf.nn.max_pool(relu1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')  

        #第二层
        conv2=tf.nn.bias_add(tf.nn.conv2d(pool1, self.weights['conv2'], strides=[1, 1, 1, 1], padding='VALID'),
                             self.biases['conv2'])
        relu2= tf.nn.relu(conv2)
        pool2=tf.nn.max_pool(relu2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')  

        # 第三层
        conv3=tf.nn.bias_add(tf.nn.conv2d(pool2, self.weights['conv3'], strides=[1, 1, 1, 1], padding='VALID'),
                             self.biases['conv3'])
        relu3= tf.nn.relu(conv3)
        pool3=tf.nn.max_pool(relu3, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='VALID')  

        # 全连接层1，先把特征图转为向量
        flatten = tf.reshape(pool3, [-1, self.weights['fc1'].get_shape().as_list()[0]])
        drop1=tf.nn.dropout(flatten,0.5)
        fc1=tf.matmul(drop1, self.weights['fc1'])+self.biases['fc1']
        fc_relu1=tf.nn.relu(fc1)
        fc2=tf.matmul(fc_relu1, self.weights['fc2'])+self.biases['fc2']
        return  fc2  

    #计算softmax交叉熵损失函数
    def softmax_loss(self,predicts,labels):
        predicts=tf.nn.softmax(predicts)
        labels=tf.one_hot(labels,self.weights['fc2'].get_shape().as_list()[1])
        loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits = predicts, labels =labels))
        self.cost= loss
        return self.cost
    #梯度下降
    def optimer(self,loss,lr=0.01):
        train_optimizer = tf.train.GradientDescentOptimizer(lr).minimize(loss)  

        return train_optimizer  

def train():
    image,label=read_and_decode("./train.tfrecords")
    batch_image,batch_label=get_batch(image,label,batch_size=30,crop_size=39)
   #建立网络,训练所用
    net=network()
    inf=net.buildnet(batch_image)
    loss=net.softmax_loss(inf,batch_label)  #计算loss
    opti=net.optimer(loss)  #梯度下降

    init=tf.global_variables_initializer()
    with tf.Session() as session:
        with tf.device("/gpu:0"):
            session.run(init)
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(coord=coord)
            max_iter=1000
            iter=0
            if os.path.exists(os.path.join("model",'model.ckpt')) is True:
                tf.train.Saver(max_to_keep=None).restore(session, os.path.join("model",'model.ckpt'))
            while iter<max_iter:
                loss_np,_,label_np,image_np,inf_np=session.run([loss,opti,batch_image,batch_label,inf])
                if iter%50==0:
                    print 'trainloss:',loss_np
                iter+=1
            coord.request_stop()#queue需要关闭，否则报错
            coord.join(threads)
if __name__ == '__main__':
    train()

结果如下：

Total memory: 10.91GiB
Free memory: 10.16GiB
2018-02-02 10:13:24.462286: I tensorflow/core/common_runtime/gpu/gpu_device.cc:961] DMA: 0
2018-02-02 10:13:24.462294: I tensorflow/core/common_runtime/gpu/gpu_device.cc:971] 0:   Y
2018-02-02 10:13:24.462303: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1030] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:03:00.0)
trainloss: 0.745739
trainloss: 0.330364
trainloss: 0.317668
trainloss: 0.314964
trainloss: 0.314613
trainloss: 0.314483
trainloss: 0.314132
trainloss: 0.313661