Tensorflow模型压缩及调用

1. 使用saver将学习的checkpoint导出成pb

import os
os.environ['CUDA_VISIBLE_DEVICES']='0'  #设置可用的GPU
import tensorflow as tf
# model的路径，checkpoint,一般包含三个文件.meta,.index和model.ckpt-10000.data.xxxxx-of-xxxxx
model_path  = "path to /model.ckpt-0000" # 只需要写到.data之前三个文件一样的内容即可
tf.reset_default_graph()
saver = tf.train.Saver()
with tf.Session() as sess:
    saver.restore(sess, model_path)
    # 保存图到.pb文件
    tf.train.write_graph(sess.graph_def, '保存model的文件夹路径', 'modelname.pb')

2. 使用freeze_graph将权重也保存到pb文件之中

from tensorflow.python.tools import freeze_graph

freeze_graph.freeze_graph(path/to/model.pb', # 之前保存出来的pb的 路径
     '', # input_saver_def_path
    False,# input_binary, true代表pb已经读入内存是个stream
    model_path,# checkpoint路径，写法为"path to /model.ckpt-0000"，用于读取checkpoint中保存的权重值
    'out',# 输出节点名，与网络预设对应
    'save/restore_all', # restore_op_name
    'input', # 输入节点名，，与网络预设对应
    'output_model.pb', # 输出的pb的路径
    False, # clear_devices
   '')

3. 使用pb文件，

    config = tf.ConfigProto(device_count={'GPU': 1})
    config.gpu_options.allow_growth = True
    with tf.Graph().as_default():
        output_graph_def = tf.GraphDef()
        pbpath = 'model.pb' # pb路径
        with open(pbpath, "rb") as f:
            output_graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(output_graph_def, name="")

        with tf.Session(config=config) as sess:
            init = tf.global_variables_initializer()
            sess.run(init)
            images = sess.graph.get_tensor_by_name("input:0")
            pred = sess.graph.get_tensor_by_name("out:0")
            data= getdata()# 获取要输入的数据
            pred = sess.run(pred, feed_dict={images: data})

4. 压缩pb文件

需要用到tensorflow下的工具graph_transforms,首先需要下载最新的tensorflow源码。
然后安装bazel进行工具编译。
编译完之后,使用命令行进行压缩。
以上操作建议在ubuntu操作系统下进行，windows系统下bazel的安装、配置和工具的编译非常麻烦，极易出错。
如果一定要在windows下编译建议编译1.8版使用cmake进行编译。
压缩后的pb文件调用方法同见3，graph_transforms命令行如下：(具体可以参考谷歌相关文档)

transform_graph --in_graph="/path/to/input.pb" --out_graph="path/to/save.pb"  --inputs="input_node_name"  --outputs="out_node_name" --transforms="
              strip_unused_nodes   \
              fold_batch_norms \
              fold_old_batch_norms quantize_nodes  \
               sort_by_execution_order"