Skip to main content
 首页 » 编程设计

python之如何正确减小 tensorflow 保存模型的大小

2024年12月31日33zhenyulu

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

请您参考如下方法:

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

  • 卡住:将保存在 SavedModel 的检查点文件中的变量转换为直接保存在模型图中的常量。这减小了模型的整体尺寸。
  • 修剪:去除预测路径中未使用的节点和图的输出,合并重复节点,以及清理其他节点操作,如摘要、身份等。
  • 常量折叠:在模型中查找总是计算为常量表达式的任何子图,并将它们替换为这些常量。
    折叠批规范:将批归一化中引入的乘法折叠到前一层的权重乘法中。
  • 量化:将权重从浮点数转换为较低的精度,例如 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)