概述

以前自己都利用別人搭好的工程，修改過來用，很少把模型搭建、匯出模型、載入模型執行走一遍，搞了一遍才知道這個事情也不是那麼簡單的。

搭建模型和匯出模型

參考《TensorFlow固化模型》，匯出固化的模型有兩種方式.

方式1：匯出pb圖結構和ckpt檔案，然後用 freeze_graph 工具凍結生成一個pb（包含結構和引數）

在我的程式碼裡測試了生成pb圖結構和ckpt檔案，但是沒接著往下走，感覺有點麻煩。我用的是第二種方法。

注意我這裡只在最後儲存了一次ckpt，實際應該在訓練中每隔一段時間就儲存一次的。

 saver = tf.train.Saver(max_to_keep=5)
 #tf.train.write_graph(session.graph_def,FLAGS.model_dir,"nn_model.pbtxt",as_text=True)
 
 with tf.Session() as sess:
 sess.run(tf.global_variables_initializer())

 max_step = 2000
 for i in range(max_step):
 batch = mnist.train.next_batch(50)
 if i % 100 == 0:
 train_accuracy = accuracy.eval(feed_dict={
  x: batch[0],y_: batch[1],keep_prob: 1.0})
 print('step %d,training accuracy %g' % (i,train_accuracy))
 train_step.run(feed_dict={x: batch[0],keep_prob: 0.5})
 
 print('test accuracy %g' % accuracy.eval(feed_dict={
 x: mnist.test.images,y_: mnist.test.labels,keep_prob: 1.0}))
 
 # 儲存pb和ckpt
 print('save pb file and ckpt file')
 tf.train.write_graph(sess.graph_def,graph_location,"graph.pb",as_text=False)
 checkpoint_path = os.path.join(graph_location,"model.ckpt")
 saver.save(sess,checkpoint_path,global_step=max_step)

方式2：convert_variables_to_constants

我實際使用的就是這種方法。

看名字也知道，就是把變數轉化為常量儲存，這樣就可以愉快的載入使用了。

注意這裡需要指明儲存的輸出節點，我的輸出節點為'out/fc2'（我猜測會根據輸出節點的依賴推斷哪些部分是訓練用到的，推理時用不到）。關於輸出節點的名字是有規律的，其中out是一個name_scope名字，fc2是op節點的名字。

 with tf.Session() as sess:
 sess.run(tf.global_variables_initializer())

 max_step = 2000
 for i in range(max_step):
 batch = mnist.train.next_batch(50)
 if i % 100 == 0:
 train_accuracy = accuracy.eval(feed_dict={
  x: batch[0],keep_prob: 1.0}))

 print('save frozen file')
 pb_path = os.path.join(graph_location,'frozen_graph.pb')
 print('pb_path:{}'.format(pb_path))

 # 固化模型
 output_graph_def = convert_variables_to_constants(sess,sess.graph_def,output_node_names=['out/fc2'])
 with tf.gfile.FastGFile(pb_path,mode='wb') as f:
 f.write(output_graph_def.SerializeToString())

上述程式碼會在訓練後把訓練好的計算圖和引數儲存到frozen_graph.pb檔案。後續就可以用這個模型來測試圖片了。

方式2的完整訓練和儲存模型程式碼

主要看main函式就行。另外注意deepnn函式最後節點的名字。

"""A deep MNIST classifier using convolutional layers.

See extensive documentation at
https://www.tensorflow.org/get_started/mnist/pros
"""
# Disable linter warnings to maintain consistency with tutorial.
# pylint: disable=invalid-name
# pylint: disable=g-bad-import-order

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import sys
import tempfile
import os

from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.python.framework.graph_util import convert_variables_to_constants

import tensorflow as tf
FLAGS = None

def deepnn(x):
 """deepnn builds the graph for a deep net for classifying digits.

 Args:
 x: an input tensor with the dimensions (N_examples,784),where 784 is the
 number of pixels in a standard MNIST image.

 Returns:
 A tuple (y,keep_prob). y is a tensor of shape (N_examples,10),with values
 equal to the logits of classifying the digit into one of 10 classes (the
 digits 0-9). keep_prob is a scalar placeholder for the probability of
 dropout.
 """
 # Reshape to use within a convolutional neural net.
 # Last dimension is for "features" - there is only one here,since images are
 # grayscale -- it would be 3 for an RGB image,4 for RGBA,etc.
 with tf.name_scope('reshape'):
 x_image = tf.reshape(x,[-1,28,1])

 # First convolutional layer - maps one grayscale image to 32 feature maps.
 with tf.name_scope('conv1'):
 W_conv1 = weight_variable([5,5,1,32])
 b_conv1 = bias_variable([32])
 h_conv1 = tf.nn.relu(conv2d(x_image,W_conv1) + b_conv1)

 # Pooling layer - downsamples by 2X.
 with tf.name_scope('pool1'):
 h_pool1 = max_pool_2x2(h_conv1)

 # Second convolutional layer -- maps 32 feature maps to 64.
 with tf.name_scope('conv2'):
 W_conv2 = weight_variable([5,32,64])
 b_conv2 = bias_variable([64])
 h_conv2 = tf.nn.relu(conv2d(h_pool1,W_conv2) + b_conv2)

 # Second pooling layer.
 with tf.name_scope('pool2'):
 h_pool2 = max_pool_2x2(h_conv2)

 # Fully connected layer 1 -- after 2 round of downsampling,our 28x28 image
 # is down to 7x7x64 feature maps -- maps this to 1024 features.
 with tf.name_scope('fc1'):
 W_fc1 = weight_variable([7 * 7 * 64,1024])
 b_fc1 = bias_variable([1024])

 h_pool2_flat = tf.reshape(h_pool2,7 * 7 * 64])
 h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat,W_fc1) + b_fc1)

 # Dropout - controls the complexity of the model,prevents co-adaptation of
 # features.
 with tf.name_scope('dropout'):
 keep_prob = tf.placeholder(tf.float32,name='ratio')
 h_fc1_drop = tf.nn.dropout(h_fc1,keep_prob)

 # Map the 1024 features to 10 classes,one for each digit
 with tf.name_scope('out'):
 W_fc2 = weight_variable([1024,10])
 b_fc2 = bias_variable([10])

 y_conv = tf.add(tf.matmul(h_fc1_drop,W_fc2),b_fc2,name='fc2')
 return y_conv,keep_prob

def conv2d(x,W):
 """conv2d returns a 2d convolution layer with full stride."""
 return tf.nn.conv2d(x,W,strides=[1,1],padding='SAME')

def max_pool_2x2(x):
 """max_pool_2x2 downsamples a feature map by 2X."""
 return tf.nn.max_pool(x,ksize=[1,2,padding='SAME')

def weight_variable(shape):
 """weight_variable generates a weight variable of a given shape."""
 initial = tf.truncated_normal(shape,stddev=0.1)
 return tf.Variable(initial)

def bias_variable(shape):
 """bias_variable generates a bias variable of a given shape."""
 initial = tf.constant(0.1,shape=shape)
 return tf.Variable(initial)

def main(_):
 # Import data
 mnist = input_data.read_data_sets(FLAGS.data_dir)

 # Create the model
 with tf.name_scope('input'):
 x = tf.placeholder(tf.float32,[None,784],name='x')

 # Define loss and optimizer
 y_ = tf.placeholder(tf.int64,[None])

 # Build the graph for the deep net
 y_conv,keep_prob = deepnn(x)

 with tf.name_scope('loss'):
 cross_entropy = tf.losses.sparse_softmax_cross_entropy(
 labels=y_,logits=y_conv)
 cross_entropy = tf.reduce_mean(cross_entropy)

 with tf.name_scope('adam_optimizer'):
 train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

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

 graph_location = './model'
 print('Saving graph to: %s' % graph_location)
 train_writer = tf.summary.FileWriter(graph_location)
 train_writer.add_graph(tf.get_default_graph())

 saver = tf.train.Saver(max_to_keep=5)
 #tf.train.write_graph(session.graph_def,keep_prob: 1.0}))
 
 # save pb file and ckpt file
 #print('save pb file and ckpt file')
 #tf.train.write_graph(sess.graph_def,as_text=False)
 #checkpoint_path = os.path.join(graph_location,"model.ckpt")
 #saver.save(sess,global_step=max_step)

 print('save frozen file')
 pb_path = os.path.join(graph_location,'frozen_graph.pb')
 print('pb_path:{}'.format(pb_path))

 output_graph_def = convert_variables_to_constants(sess,mode='wb') as f:
 f.write(output_graph_def.SerializeToString())

if __name__ == '__main__':
 parser = argparse.ArgumentParser()
 parser.add_argument('--data_dir',type=str,default='./data',help='Directory for storing input data')
 FLAGS,unparsed = parser.parse_known_args()
 tf.app.run(main=main,argv=[sys.argv[0]] + unparsed)
 

載入模型進行推理

上一節已經訓練並匯出了frozen_graph.pb。

這一節把它執行起來。

載入模型

下方的程式碼用來載入模型。推理時計算圖裡共兩個placeholder需要填充資料，一個是圖片（這不廢話嗎），一個是drouout_ratio，drouout_ratio用一個常量作為輸入，後續就只需要輸入圖片了。

graph_location = './model'
pb_path = os.path.join(graph_location,'frozen_graph.pb')
print('pb_path:{}'.format(pb_path))

newInput_X = tf.placeholder(tf.float32,name="X")
drouout_ratio = tf.constant(1.,name="drouout")
with open(pb_path,'rb') as f:
 graph_def = tf.GraphDef()
 graph_def.ParseFromString(f.read())

 output = tf.import_graph_def(graph_def,input_map={'input/x:0': newInput_X,'dropout/ratio:0':drouout_ratio},return_elements=['out/fc2:0'])

input_map引數並不是必須的。如果不用input_map，可以在run之前用tf.get_default_graph().get_tensor_by_name獲取tensor的控制代碼。但是我覺得這種方法不是很友好，我這裡沒用這種方法。

注意input_map裡的tensor名字是和搭計算圖時的name_scope和op名字有關的，而且後面要補一個‘：0'(這點我還沒細究)。

同時要注意，newInput_X的形狀是[None,784]，第一維是batch大小，推理時和訓練要一致。

（我用的是mnist圖片，訓練時每個bacth的形狀是[batchsize,每個圖片是28x28）

執行模型

我是一張張圖片單獨測試的，執行模型之前先把圖片變為[1,784]，以符合newInput_X的維數。

with tf.Session( ) as sess:
 file_list = os.listdir(test_image_dir)
 
 # 遍歷檔案
 for file in file_list:
 full_path = os.path.join(test_image_dir,file)
 print('full_path:{}'.format(full_path))
 
 # 只要黑白的，大小控制在(28,28)
 img = cv2.imread(full_path,cv2.IMREAD_GRAYSCALE )
 res_img = cv2.resize(img,(28,28),interpolation=cv2.INTER_CUBIC) 
 # 變成長784的一維資料
 new_img = res_img.reshape((784))
 
 # 增加一個維度，變為 [1,784]
 image_np_expanded = np.expand_dims(new_img,axis=0)
 image_np_expanded.astype('float32') # 型別也要滿足要求
 print('image_np_expanded shape:{}'.format(image_np_expanded.shape))
 
 # 注意注意，我要呼叫模型了
 result = sess.run(output,feed_dict={newInput_X: image_np_expanded})
 
 # 出來的結果去掉沒用的維度
 result = np.squeeze(result)
 print('result:{}'.format(result))
 #print('result:{}'.format(sess.run(output,feed_dict={newInput_X: image_np_expanded})))
 
 # 輸出結果是長度為10（對應0-9）的一維資料，最大值的下標就是預測的數字
 print('result:{}'.format( (np.where(result==np.max(result)))[0][0] ))

注意模型的輸出是一個長度為10的一維陣列，也就是計算圖裡全連線的輸出。這裡沒有softmax，只要取最大值的下標即可得到結果。

輸出結果：

full_path:./test_images/97_7.jpg
image_np_expanded shape:(1,784)
result:[-1340.37145996 -283.72436523 1305.03320312 437.6053772 -413.69961548
 -1218.08166504 -1004.83807373 1953.33984375 42.00457001 -504.43829346]
result:7

full_path:./test_images/98_6.jpg
image_np_expanded shape:(1,784)
result:[ 567.4041748 -550.20904541 623.83496094 -1152.56884766 -217.92695618
 1033.45239258 2496.44750977 -1139.23620605 -5.64091825 -615.28491211]
result:6

full_path:./test_images/99_9.jpg
image_np_expanded shape:(1,784)
result:[ -532.26409912 -1429.47277832 -368.58096313 505.82876587 358.42163086
 -317.48199463 -1108.6829834 1198.08752441 289.12286377 3083.52539062]
result:9

載入模型進行推理的完整程式碼

import sys
import os
import cv2
import numpy as np
import tensorflow as tf
test_image_dir = './test_images/'

graph_location = './model'
pb_path = os.path.join(graph_location,'rb') as f:
 graph_def = tf.GraphDef()
 graph_def.ParseFromString(f.read())
 #output = tf.import_graph_def(graph_def)
 output = tf.import_graph_def(graph_def,return_elements=['out/fc2:0'])

with tf.Session( ) as sess:
 file_list = os.listdir(test_image_dir)
 
 # 遍歷檔案
 for file in file_list:
 full_path = os.path.join(test_image_dir,feed_dict={newInput_X: image_np_expanded})))
 
 # 輸出結果是長度為10（對應0-9）的一維資料，最大值的下標就是預測的數字
 print('result:{}'.format( (np.where(result==np.max(result)))[0][0] ))
 

以上這篇tensorflow 20:搭網路,匯出模型,執行模型的例項就是小編分享給大家的全部內容了，希望能給大家一個參考，也希望大家多多支援我們。