如何适当减少张量流保存模型的大小?

时间:2018-08-21 22:04:39

标签: python tensorflow tensorflow-serving

我有一个检查点形式的tensorflow预训练模型,我打算通过将模型转换为not null形式来部署该模型以进行服务。保存的模型的大小太大。 (保存的模型中的“ variables.data-00000-of-0001”文件超过100 MB。)我用谷歌搜索了如何减小变量的大小,但找不到很好的答案。您能帮助我了解如何减少tensorflow保存模型中变量的大小吗?展示一个简单的例子将是很棒的。谢谢!

1 个答案:

答案 0 :(得分:0)

我们可以使用以下提到的方法减小Tensorflow模型的大小:

  1. 冻结:将SavedModel的检查点文件中存储的变量转换为直接存储在模型图中的常量。这样可以减小模型的整体大小。

  2. 修剪:在预测路径和图形输出中剥离未使用的节点,合并重复的节点,以及清除其他节点操作,例如摘要,标识等。

  3. 恒定折叠:在模型中查找始终求值为常量表达式的任何子图,然后将其替换为这些常量。 折叠批处理规范:将批量归一化中引入的乘法折叠为上一层的权重乘法。

  4. 量化:将权重从浮点转换为较低的精度,例如16或8位。

冻结图形的代码如下:

from tensorflow.python.tools import freeze_graph

output_graph_filename = os.path.join(saved_model_dir, output_filename)
initializer_nodes = ''

freeze_graph.freeze_graph(input_saved_model_dir=saved_model_dir,
      output_graph=output_graph_filename,
      saved_model_tags = tag_constants.SERVING,
      output_node_names=output_node_names,initializer_nodes=initializer_nodes,
      input_graph=None, input_saver=False, input_binary=False, 
      input_checkpoint=None, restore_op_name=None, filename_tensor_name=None,
      clear_devices=False, input_meta_graph=False)

修剪和恒定折叠的代码如下:

from tensorflow.tools.graph_transforms import TransformGraph

def get_graph_def_from_file(graph_filepath):
  with ops.Graph().as_default():
    with tf.gfile.GFile(graph_filepath, 'rb') as f:
      graph_def = tf.GraphDef()
      graph_def.ParseFromString(f.read())
      return graph_def

def optimize_graph(model_dir, graph_filename, transforms, output_node):
  input_names = []
  output_names = [output_node]
  if graph_filename is None:
    graph_def = get_graph_def_from_saved_model(model_dir)
  else:
    graph_def = get_graph_def_from_file(os.path.join(model_dir, 
         graph_filename))
  optimized_graph_def = TransformGraph(graph_def, input_names,      
      output_names, transforms)
  tf.train.write_graph(optimized_graph_def, logdir=model_dir, as_text=False, 
     name='optimized_model.pb')
  print('Graph optimized!')

我们通过传递所需的优化列表来调用模型上的代码,如下所示:

transforms = ['remove_nodes(op=Identity)', 'merge_duplicate_nodes',
 'strip_unused_nodes','fold_constants(ignore_errors=true)',
 'fold_batch_norms']

optimize_graph(saved_model_dir, "frozen_model.pb" , transforms, 'head/predictions/class_ids')

量化代码如下:

transforms = ['quantize_nodes', 'quantize_weights',]
optimize_graph(saved_model_dir, None, transforms, 'head/predictions/class_ids')

应用优化后,我们需要将“优化图”转换回GraphDef。其代码如下所示:

def convert_graph_def_to_saved_model(export_dir, graph_filepath):
  if tf.gfile.Exists(export_dir):
    tf.gfile.DeleteRecursively(export_dir)
  graph_def = get_graph_def_from_file(graph_filepath)
  with tf.Session(graph=tf.Graph()) as session:
    tf.import_graph_def(graph_def, name='')
    tf.saved_model.simple_save(
        session,
        export_dir,
        inputs={
            node.name: session.graph.get_tensor_by_name(
                '{}:0'.format(node.name))
            for node in graph_def.node if node.op=='Placeholder'},
        outputs={'class_ids': session.graph.get_tensor_by_name(
            'head/predictions/class_ids:0')}
    )
    print('Optimized graph converted to SavedModel!')

示例代码如下所示:

optimized_export_dir = os.path.join(export_dir, 'optimized')
optimized_filepath = os.path.join(saved_model_dir, 'optimized_model.pb')
convert_graph_def_to_saved_model(optimized_export_dir, optimized_filepath)