在Tensorflow 1.x中优化模型

Tensorflow虽然在研究环境中有所失落,但在实际开发中仍然很受欢迎TF的优势之一就是能够优化模型以在资源受限的环境中进行部署。有一些特殊的框架:用于移动设备的Tensorflow LiteTensorflow Serving用于工业。在网络上(甚至在Habré上)有足够的使用它们的教程。在本文中,我们收集了在不使用这些框架的情况下优化模型的经验。我们将研究完成手头任务的一些方法和库,描述如何节省磁盘空间和RAM,每种方法的优缺点以及遇到的一些意外影响。



我们在什么条件下工作



NLP的经典任务之一是对短文本进行主题分类。分类器由许多不同的体系结构表示,从诸如SVC的经典方法到诸如BERT及其派生类的转换器体系结构。我们将研究CNN-卷积模型。



对我们来说,一个重要的限制是需要在没有GPU的机器上训练和使用模型(作为产品的一部分)。这主要影响学习和推理的速度。



另一个条件是,要对分类模型进行训练并以几件为一组使用。一组模型,甚至是简单的模型,都可以使用很多资源,尤其是RAM。我们使用自己的服务模型解决方案,但是,如果您需要使用模型集进行操作,请查看Tensorflow服务



我们面临着在TF版本1.x上优化模型的需求,现在正式认为该版本已过时。对于TF 2.x,讨论的许多技术都不相关或已集成到标准API中,因此优化过程非常简单。



首先让我们看一下模型的结构。



TF模型如何运作



考虑所谓的浅CNN-具有一个卷积层和几个过滤器的网络。该模型已经可以很好地用于矢量单词表示形式的文本分类。





为简单起见,我们将使用固定的预训练向量集维v x k,其中v是字典的大小,k是嵌入的维。



:



  • Embedding-, .
  • w x k. , (1, 1, 2, 3) 4 , 1 , 2 3 , .
  • Max-pooling .
  • , dropout- softmax- .


Adam, .



: .



, , 128 c w = 2 k = 300 () [filter_height, filter_width, in_channels, output_channels] — , 2*300*1*128 = 76800 float32, , 76800*(32/8) = 307200 .



? ( 220 . ) 300 265 . , .



TF . ( ), , , — ( ), . (). :



计算图





. , : SavedModel. , .



Checkpoint



, Saver API:



saver = tf.train.Saver(save_relative_paths=True)
ckpt_filepath = saver.save(sess, "cnn.ckpt"), global_step=0)


global_step , , — cnn-ckpt-0.



<model_path>/cnn_ckpt :





checkpoint — . , TF . , .



.data , . , — 800 . , (≈265 ). ( ). , .



.index .



.meta — , (, , ), GraphDef, . , . — .meta , ? , TF - embedding-. , , , , , . , , :



with tf.Session() as sess:
   saver = tf.train.import_meta_graph('models/ckpt_model/cnn_ckpt/cnn.ckpt-0.meta')  # load meta
   for n in tf.get_default_graph().as_graph_def().node:
       print(n.name, n['attr'].shape)


.



SavedModel



, . . API tf.saved_model. tf.saved_model, TF- (TFLite, TensorFlow.js, TensorFlow Serving, TensorFlow Hub).



:





saved_model.pb, , , .meta , (, ), API, ( CLI, ).



SavedModel — , . “” . , , - — , .





, CNN-, TF 1.x, . .



, 1 , :





  1. . , , ( tools.optimize_for_inference ).


  2. . , , — , tf.trainable_variables().


  3. , . , (. BERT).


  4. . , . .




, , . , forward pass, . , . 1 265 .



TF 1.x , .



( ) GraphDef:



graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()


. : tf.python.tools.freeze_graph tf.graph_util.convert_variables_to_constants. ( ) (, ['output/predictions']), , , . .



output_graph_def = graph_util.convert_variables_to_constants(self.sess, input_graph_def, output_node_names)


, .

freeze_graph() ( , , ). graph_util.convert_variables_to_constants() :



with tf.io.gfile.GFile('graph.pb', 'wb') as f:
    f.write(output_graph_def.SerializeToString())


266 , :



#  GraphDef  

with tf.io.gfile.GFile(graph_filepath, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

with tf.Graph().as_default() as graph:
    #   
    self.input_x = tf.placeholder(tf.int32, [None, self.properties.max_len], name="input_x")
    self.dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
    #        graph_def
    input_map = {'input_x': self.input_x, 'dropout_keep_prob': self.dropout_keep_prob}
    tf.import_graph_def(graph_def, input_map)


, import:



predictions = graph.get_tensor_by_name('import/output/predictions:0')


:



feed_dict = {self.input_x: encode_sentence(sentence), self.dropout_keep_prob: 1.}
sess.run(self.predictions, feed_dict)


, :



  1. . , sess.run(...). , CPU 20 ms, ~2700 ms. , . SavedModel .
  2. RAM. RAM, . ~265 , . , TF GraphDef .
  3. – RAM TF . 1.15, TF 1.x, 118 MiB, 1.14 – 3 MiB.




, . ? / TF- tf.train.Saver. , , , :



  • MetaGraph


tf.train.Saver . , :



saver = tf.train.Saver(var_list=tf.trainable_variables())


MetaGraph . , meta . MetaGraph save:



ckpt_filepath = saver.save(self.sess, filepath, write_meta_graph=False)


1014 M 265 M ( , ).





Pruning — , , . , .



, TF 1.x:



  • Grappler: c tensorflow
  • Pruning API: google-research
  • Graph Transform Tool:


, — tensorflow, Grappler. Grappler . , set_experimental_options. , zip . , zip , . Grappler .



google-research mask threshold, . . , , mask threshold, , , . .



Grappler, . : ? , ? , 0.99 . , mc, hex :



hex-



, , . . -, . -, , , , . , .



CNN. .





, . Graph transform tool.



quantize_weights 8 . , 8- . , , - .



quantize_nodes 8- . .



, - . quantize_weights - , 4 .



, , TensorFlow Lite, .





— , . 64 (32) , .



RAM Ubuntu ( numpy int64) . 220 , int32, int16. .





tf-. float16. , , ( 10%), ( 10 ). , , epsilon learning_rate . , , .



RAM



, . , .





, . . .



QA-



Q: -, - ?



A: , . word2vec. ( , , min count, learning rate), 220 ( — 265 MB) CNN, 439 (510 MB).



- , , , - . , ( ). , . YouTokenToMe, , , .. , .., . . , , , . 30 (37 MB) , 3.7 CPU 2.6 GPU. ( ), OOV-.



Q: , , ?



A: , .



:



1. :



with tf.gfile.GFile(path_to_pb, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
    tf.import_graph_def(graph_def, name='')
    return graph


2. "" :



sess.run(restored_variable_names)


3. , .

4. , , :



tf.Variable(tensors_to_restore["output/W:0"], name="W")


, .



, , .



我们没有尝试重新训练由上述方法压缩的模型,但是从理论上讲,这应该没有任何问题。



问:还有其他没有考虑过的减少优化的方法吗?



答:我们有一些我们从未意识到的想法。首先,常量折叠是图节点子集的``折叠'',即对图的部分依赖于输入数据的部分的值进行预计算。其次,在我们的模型中,应用嵌入修剪似乎是一个很好的解决方案。




All Articles