tf-CNN（三）Tensorflow將訓練模型保存到本地

訓練過程在上一篇博客裏https://blog.csdn.net/xue_csdn/article/details/105128094

首先在代碼的相同路徑下新建一個文件夾save_model。模型會被保存進來。

我這兒就直接把網絡結構和模型的訓練和保存寫在同一代碼裏，不再互相調用了。

save_model.py

"""
###訓練貓狗分類模型
###保存生成的模型

"""

#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
import numpy as np
import read_tfrecords

epoch = 15
batch_size = 20

def one_hot(labels, Label_class):
    one_hot_label = np.array([[int(i==int(labels[j])) for i in range(Label_class)] for j in range(len(labels))])
    return one_hot_label

#初始化權值
def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev = 0.02)
    return tf.Variable(initial)

#初始化偏置
def bias_variable(shape):
    initial = tf.constant(0.0 , shape=shape)
    return tf.Variable(initial)

#卷積層
def conv2d(x,W):
    return tf.nn.conv2d(x,W,strides = [1,1,1,1],padding = 'SAME')

#池化層
def max_pool_4x4(x):
    return tf.nn.max_pool(x,ksize=[1,4,4,1],strides=[1,4,4,1],padding='SAME')

#佔位符
x = tf.placeholder(tf.float32, shape=[batch_size, 128,128,3],name = 'x')
y_ = tf.placeholder(tf.float32, shape=[batch_size, 2],name = 'y_')

#1 卷積1和池化1
W_conv1 = weight_variable([5,5,3,32])
b_conv1 = bias_variable([32])
h_conv1 = tf.nn.relu(conv2d(x,W_conv1)+b_conv1)
h_pool1 = max_pool_4x4(h_conv1)

#2 卷積2和池化2
W_conv2 = weight_variable([5,5,32,64])
b_conv2 = bias_variable([64])
h_conv2 = tf.nn.relu(conv2d(h_pool1,W_conv2)+b_conv2)
h_pool2 = max_pool_4x4(h_conv2)

#全連接，用1個MLP處理
reshape = tf.reshape(h_pool2,[batch_size, -1])
dim = reshape.get_shape()[1].value
W_fc1 = weight_variable([dim, 1024])
b_fc1 = bias_variable([1024])
h_fc1 = tf.nn.relu(tf.matmul(reshape, W_fc1)+b_fc1)

#dropout
keep_prob = tf.placeholder(tf.float32)
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

W_fc2 = weight_variable([1024,2])
b_fc2 = bias_variable([2])
y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2)+b_fc2)

############
tf.add_to_collection('network-output',y_conv)


#損失函數及優化算法
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y_conv),reduction_indices = [1]))
train_step = tf.train.AdamOptimizer(0.001).minimize(cross_entropy)

correct_prediction = tf.equal(tf.argmax(y_conv,1),tf.argmax(y_,1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))

img,label = read_data2.read_and_decode("dog_and_cat_train.tfrecords")
img_test, label_test = read_data2.read_and_decode(("dog_and_cat_test.tfrecords"))

#使用shuffle_batch可以隨機打亂輸入
img_batch, label_batch = tf.train.shuffle_batch([img,label],
                                                batch_size=batch_size,capacity=2000,
                                                min_after_dequeue=1000)
img_test,label_test = tf.train.shuffle_batch([img_test,label_test],
                                             batch_size=batch_size,capacity=2000,
                                             min_after_dequeue=1000)

init = tf.initialize_all_variables()
t_vars = tf.trainable_variables()
print(t_vars)

with tf.Session() as sess:
    sess.run(init)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess,coord=coord)
    batch_idxs = int(2314/batch_size)
    for i in range(epoch):
        for j in range(batch_idxs):
            val,l = sess.run([img_batch, label_batch])
            l = one_hot(l,2)
            _, acc = sess.run([train_step, accuracy], feed_dict={x: val, y_: l, keep_prob: 0.5})
            print("Epoch:[%4d] [%4d/%4d], accuracy:[%.8f]" % (i, j, batch_idxs, acc))

    val,l = sess.run([img_test,label_test])
    l=one_hot(l,2)
    print(l)
    y,acc = sess.run([y_conv,accuracy],feed_dict={x:val,y_:l,keep_prob:1})
    print(y)
    print("test accuracy:[%.8f]" % (acc))

    coord.request_stop()
    coord.join(threads)

    ####模型保存
    saver=tf.train.Saver()
    saver.save(sess,'./save_model/dog_cat_model.ckpt')
    sess.close()

運行之後，會在save_model文件夾下出現幾個文件

“checkpoint”：文件僅用於告知某些TF函數，這是最新的檢查點文件。
.ckpt-meta：包含元圖，即計算圖的結構，沒有變量的值（基本上可以在tensorboard / graph中看到）。
.ckpt-data：包含所有變量的值，沒有結構。
.ckpt-index：可能是內部需要的某種索引來正確映射前兩個文件，它通常不是必需的
可以只用 .ckpt-meta 和恢復一個模型 .ckpt-data

 

模型保存之後，可以調用.meta來測試圖片。這一步部分還沒琢磨明白（如果有大神還望賜教！！！），而且還需要在訓練模型的代碼中標明輸入、輸出的接口，所以代碼後期仍需修改。