TensorFlow 笔记01-mnist 的两种 tensorFlow 实现，保存和加载模型的三种方法

▶ 环境：DisplayDriver-430.50 + cuda-10.0.130_with-DisplayDriver-410.48 + cudnn-10.0-7.6.4.38 + TensorRT-5.1.5.0_attachedTo-cuda-10.0_cudnn_7.5 + Anaconda3-2019.07 + tensorflow-gpu-1.13.1

● 代码，tf 单隐层 + softmax 神经网络，需要把 input_data.py 放到执行目录下。分别使用 tf.saved_model.builder 和 tf.train.Saver 来保存模型，重新加载后进行推理。使用固定随机数种子产生可重复正确率，107 -> 0.91150

import os
import numpy as np
import tensorflow as tf

import input_data                                               # 要把 input_data.py 放到执行目录

#os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'                       # 仅输出警告信息
tempPath = "tempFile/"
tf.random.set_random_seed(107)

mnist = input_data.read_data_sets('MNIST_data/', one_hot=True)  # 读取 mnist 数据，指定one_hot 编码

# 建立模型
x = tf.placeholder(tf.float32, [None, 784], name='x')           # 占位符，数据类型（也可用 "float"），尺寸（None 表示自动计算），节点名称（可以没有）
W = tf.Variable(tf.zeros([784, 10]), name='W')                  # 变量，自动加入 tf.trainable_variables（可训练变量）集合，除非指定 trainable = False
b = tf.Variable(tf.ones([10]), name='b')                        # tf.zeros() 全零张量
y = tf.nn.softmax(tf.matmul(x, W) + b, name='y')                # tf.nn.softmax() softmax 层
y_ = tf.placeholder("float", [None, 10], name='y_')
res = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))                 # tf.equal() 逐元相等判断，tf.argmax() 取最大元素下标
acc = tf.reduce_mean(tf.cast(res, "float"), name = 'acc')       # tf.reduce_mean() 求平均值，tf.cast() 强制类型转换

cross_entropy = -tf.reduce_sum(y_*tf.log(y))                    # 定义交叉熵
train_step = tf.train.GradientDescentOptimizer(0.01).minimize(cross_entropy)# 定义单步训练，采用梯度下降法，学习率 0.01

sess = tf.Session()                                             # 开启一个会话，可用 sess = InteractiveSession()
sess.run(tf.initialize_all_variables())                         # 初始化变量，新版本用 tf.global_variables_initializer().run() 或 sess.run( tf.global_...() )

for i in range(1000):
  xSample, ySample = mnist.train.next_batch(100)                # 随机取 batchSize 的元素作为一次输入
  sess.run(train_step, feed_dict={x: xSample, y_: ySample})     # 执行单步训练，自动更新权重

print( "acc = %f"%sess.run(acc, feed_dict={x: mnist.test.images, y_: mnist.test.labels}) )

# 保存模型
builder = tf.saved_model.builder.SavedModelBuilder(tempPath+'para')  # 用 tf.saved_model.builder.SavedModelBuilder 来保存模型
builder.add_meta_graph_and_variables(sess,['y','W','b']) 
builder.save()

#saver = tf.train.Saver()                                        # 用 tf.train.Saver 来保存模型为 checkpoint 文件
#saver.save(sess,tempPath + 'model.ckpt')

sess.close()                                                    # 没有使用 with sess = tf.Sesion() 上下文的情况下要记得关闭会话

# 加载模型
sess = tf.Session()                                            
sess.run(tf.global_variables_initializer())                     

tf.saved_model.loader.load(sess, ['y','W','b'], tempPath+'para')# 加载用 SavedModelBuilder 保存的模型

#saver.restore(sess, tempPath+'model.ckpt')                     # 加载用 tf.train.Saver 保存的模型

x = sess.graph.get_tensor_by_name('x:0')                        # 恢复各张量，再次测试
y_ = sess.graph.get_tensor_by_name('y_:0')
y = sess.graph.get_tensor_by_name('y:0')
acc = sess.graph.get_tensor_by_name('acc:0')

print( "acc = %f"%sess.run(acc, feed_dict={x: mnist.test.images, y_: mnist.test.labels}) )

sess.close()

● 代码，tf 卷积神经网络，卷积 -> 池化 -> 卷积 -> 池化 -> 全连接 -> 全连接 -> softmax，将模型保存为 .pb 文件

import os
import numpy as np
import tensorflow as tf
from datetime import datetime as dt

import input_data

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tempPath = "tempFile/"

def wInit(shape, name):                                                     # 初始化 w 和 b 的函数
    return tf.Variable(tf.truncated_normal(shape, stddev=0.1), name = name) # truncated_normal() 给定范围和方差的截断标准正态分布

def bInit(shape, name):
    return tf.Variable(tf.constant(0.1, shape=shape), name = name)          # tf.constant() 赋予常量

def conv(x, W):                                                             # 卷积层
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')         # tf.nn.conv2d() 二维卷积，步长 1（第 0 和第 3 元素固定为 1），补齐原图像尺寸

def maxPool(x):                                                             # 池化层
    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')  # tf.nn.max_pool max 池化，窗口 2×2，步长 2 （第 0 和第 3 元素均固定为 1），总体尺寸缩小 1 倍

# 读数据，建图
mnist = input_data.read_data_sets('MNIST_data/', one_hot=True)

x =  tf.placeholder(tf.float32, [None, 784], name='x')
y_ = tf.placeholder(tf.float32, [None,10], name='y_')

xImage = tf.reshape(x, [-1,28,28,1], name='xImage')                         # x 调整为 ? 页 28 行 28 列 1 通道，? 自动计算
w1 = wInit([5, 5, 1, 32], name='w1')                                        # 1 层卷积窗口，5 行 5 列 1 页 32 个特征
b1 = bInit([32], name='b1')                                                 # 1 层偏置，对应 32 特征
h1 = tf.nn.relu(conv(xImage, w1) + b1)                                      # 1 层卷积，激活，池化
h1Pool = maxPool(h1)                                                        

w2 = wInit([5, 5, 32, 64], name='w2')                                       # 2 层卷积窗口，5 行 5 列 32 页 64 个特征
b2 = bInit([64], name='b2')                                                 
h2 = tf.nn.relu(conv(h1Pool, w2) + b2)                                      # 2 层卷积，激活，池化
h2Pool = maxPool(h2)                                                        

w3 = wInit([7 * 7 * 64, 1024], name='w3')                                   # 全连接层 1
b3 = bInit([1024], name='b3')                                               
h2Flat = tf.reshape(h2Pool, [-1, 7*7*64])                                   # 展平后做全连接，激活
h3 = tf.nn.relu(tf.matmul(h2Flat, w3) + b3)

w4 = wInit([1024, 10], name='w4')                                           # 全连接层 2
b4 = bInit([10], name='b4')                                                 
y=tf.nn.softmax(tf.matmul(h3, w4) + b4, name='y')                           # 输出做一个 softmax

# 训练和结果检测
cross_entropy = -tf.reduce_sum(y_*tf.math.log(y))
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)           # 这里 AdamOptimizer 效果比梯度下降好

output = tf.argmax(y,1)                                                     
resultCheck = tf.equal(tf.argmax(y,1), tf.argmax(y_,1))
acc = tf.reduce_mean(tf.cast(resultCheck, tf.float32), name='acc')

# 训练和测试
sess = tf.Session()
sess.run(tf.initialize_all_variables())
for i in range(3000):
  xSample, ySample = mnist.train.next_batch(100)
  if i%100 == 0:
    train_acc = acc.eval(session = sess, feed_dict={x:xSample, y_: ySample})
    print("%s, step %d, acc = %f"%(dt.now(), i, train_acc))
  train_step.run(session = sess, feed_dict={x: xSample, y_: ySample})

print( "%s, test acc = %f"%(dt.now(), acc.eval(session = sess, feed_dict={x: mnist.test.images, y_: mnist.test.labels})) )

# 保存模型到 pb 中
constantGraph = tf.compat.v1.graph_util.convert_variables_to_constants(sess, sess.graph_def,['y','acc'])    # 指定图和保存变量的列表
f = tf.gfile.FastGFile(tempPath+'model.pb', mode='wb')                      # gFile 打开和写入
f.write(constantGraph.SerializeToString())
f.close()
sess.close()

● 代码，tf 卷积神经网络，从 pb 文件中读取训练好的模型和参数，进行推理

import os
import numpy as np
import tensorflow as tf
from datetime import datetime as dt

import input_data

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
tempPath = "tempFile/"
mnist = input_data.read_data_sets('MNIST_data/', one_hot=True)

sess = tf.Session()
f = tensorflow.gfile.FastGFile(tempPath + "model.pb", 'rb') # 读取 .pb 文件

graph_def = tf.GraphDef()                                   # 建立默认图
graph_def.ParseFromString(f.read())                 
_ = tf.import_graph_def(graph_def, name="")                 # 加载图，name 为 scope 名

x = sess.graph.get_tensor_by_name('x:0')                    # 加载需要的节点
y_ = sess.graph.get_tensor_by_name('y_:0')
#oy = sess.graph.get_tensor_by_name('y:0')
acc = sess.graph.get_tensor_by_name('acc:0')

print( "%s, test acc = %f"%(dt.now(), sess.run(acc, feed_dict={x: mnist.test.images, y_: mnist.test.labels})) ) # 测试

f.close()
sess.close()