Tensorflow虽然在研究环境中有所失落,但在实际开发中仍然很受欢迎。TF的优势之一就是能够优化模型以在资源受限的环境中进行部署。有一些特殊的框架:用于移动设备的Tensorflow Lite和Tensorflow 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 , :
-
. , , ( tools.optimize_for_inference ). -
. , , — , tf.trainable_variables(). -
, . , (. BERT). -
. , . .
, , . , 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)
, :
- . ,
sess.run(...)
. , CPU 20 ms, ~2700 ms. , . SavedModel . - RAM. RAM, . ~265 , . , TF GraphDef .
- – 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 ( , ).
, 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 :
, , . . -, . -, , , , . , .
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")
, .
, , .
我们没有尝试重新训练由上述方法压缩的模型,但是从理论上讲,这应该没有任何问题。
问:还有其他没有考虑过的减少优化的方法吗?
答:我们有一些我们从未意识到的想法。首先,常量折叠是图节点子集的``折叠'',即对图的部分依赖于输入数据的部分的值进行预计算。其次,在我们的模型中,应用嵌入修剪似乎是一个很好的解决方案。