TensorFlow-cifar训练与测试（可对自己数据进行分类和测试）

第一部分：测试软硬件

硬件：NVIDIA-GTX1080

软件：Windows7、python3.6.5、tensorflow-gpu-1.4.0

第二部分：数据下载

数据集下载链接

第三部分：代码分步展示

第一步：导入tensorflow

import os
from PIL import Image
import numpy as np
import tensorflow as tf

第二步：定义模型函数

def model_load(x_, y_, keep_prob):
    conv1 = tf.layers.conv2d(x_, 20, 5, activation=tf.nn.relu)
    pool1 = tf.layers.max_pooling2d(conv1, [2, 2], [2, 2])
    conv2 = tf.layers.conv2d(pool1, 40, 4, activation=tf.nn.relu)
    pool2 = tf.layers.max_pooling2d(conv2, [2, 2], [2, 2])
    pool2_reshape = tf.layers.flatten(pool2)
    fc = tf.layers.dense(pool2_reshape, 400, activation=tf.nn.relu)
    fc_dropout = tf.layers.dropout(fc, keep_prob)
    y_out = tf.layers.dense(fc_dropout, 3)
    return y_out

第三步：定义数据读取函数

def read_data(data_dir):
    datas = []
    labels = []
    fpaths = []
    for fname in os.listdir(data_dir):
        fpath = os.path.join(data_dir, fname)
        fpaths.append(fpath)
        image = Image.open(fpath)
        data = np.array(image) / 255.0
        label = int(fname.split("_")[0])
        datas.append(data)
        labels.append(label)

    datas = np.array(datas)
    labels = np.array(labels)
    return fpaths, datas, labels

第四步：定义网络输入输出

#define x y size
x_ = tf.placeholder(tf.float32, [None, 32, 32, 3])
y_ = tf.placeholder(tf.int32, [None])
keep_prob = tf.placeholder(tf.float32)

#load model
y_out = model_load(x_, y_, keep_prob)

#define loss
losses = tf.nn.softmax_cross_entropy_with_logits(labels = tf.one_hot(y_, 3), logits = y_out)
mean_loss = tf.reduce_mean(losses)
optimizer = tf.train.AdamOptimizer(learning_rate=1e-2).minimize(mean_loss)

#define predict top1
predicted_labels = tf.argmax(y_out, 1)

第五步：读取cifar数据集

#read data
data_dir = "data"
model_path = "model/cifar_model"
fpaths, datas, labels = read_data(data_dir)

第六步：训练和测试（100%准确率）

#define saver and sess
saver = tf.train.Saver()
sess = tf.Session()

#train or test
train = True
if train:
    sess.run(tf.global_variables_initializer())
    for step in range(150):
        _, mean_loss_val = sess.run([optimizer, mean_loss], feed_dict={x_:datas, y_:labels, keep_prob:0.25})
        if step % 10 == 0:
            print("step = {}\tmean loss = {}".format(step, mean_loss_val))
    saver.save(sess, model_path + str(150))
else:
    labels_name = ["plane", "car", "bird"]
    saver.restore(sess, model_path)
    predicted_labels_val = sess.run(predicted_labels, feed_dict={x_:datas, y_:labels, keep_prob:0})
    img_count = 0
    for fpath in fpaths:
        print("{}\t{} =? {}".format(fpath, labels_name[labels[img_count]], labels_name[predicted_labels_val[img_count]]))
        img_count += 1

第七步：备注

若发现loss不下降，说明训练过拟合了，原因是参数初始化问题。这个没关系，删除model，重新训练一遍即可。

 

任何问题请加唯一QQ2258205918（名称samylee）！