保存Tensorflow图和变量,同时排除Adam Optimizer变量

时间:2017-12-01 23:12:43

标签: tensorflow model save adam

我使用以下方法保存张量流图和变量

 builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
 builder.add_meta_graph_and_variables(sess, ["nn"])
 builder.save(as_text=False)

我想要一个可以运行分类的最小保存文件。具体来说,我希望从保存文件中排除在训练期间创建的Adam变量。

当我在保存之前迭代tf.all_variables()返回的值时,我得到了我期望的变量:

tf.Variable 'mymodel/fully_connected/weights:0' shape=(128, 100) dtype=float32_ref

但也有两个相同的亚当副本:

tf.Variable 'mymodel/train/mymodel/fully_connected/weights/Adam:0' shape=(128, 100) dtype=float32_ref

tf.Variable 'mymodel/train/mymodel/fully_connected/weights/Adam_1:0' shape=(128, 100) dtype=float32_ref

这使我的检查点文件的大小增加了三倍,我受到目标系统的限制。保存文件仅用于分类,而不是培训,因此我不需要Adam优化器变量。

关于如何最容易将这些内容写入保存文件的任何建议?

我们赞赏任何其他有关减少保存文件大小同时仍然能够运行分类的建议。

3 个答案:

答案 0 :(得分:1)

我最终使用以下代码完成工作。

我保存了图表,重量,偏差等,如下所示:

  with tf.Session(graph=tf.Graph()) as sess:
    // Checkpoint of trained model that has Adam optimizer variables
    tf.saved_model.loader.load(sess, ["audio_nn"], FLAGS.checkpoint_dir)

    # Dump out checkpoint without Adam optimizer variables
    saver = tf.train.Saver(tf.model_variables())
    saver.save(sess, 'my-model')

幸运的是,tf.model_variables()仅返回模型所需的变量,而不返回Adam优化器变量。

并恢复为这样的分类:

 with tf.Session(graph=tf.Graph()) as sess:
    imported_meta = tf.train.import_meta_graph("my-model.meta")
    imported_meta.restore(sess, "my-model")

    features_tensor = sess.graph.get_tensor_by_name( params.INPUT_TENSOR_NAME)
    prediction_tensor = sess.graph.get_tensor_by_name( 'mymodel/prediction:0')

    ... prediction code ...

现在,检查点文件的大小约为以前的三分之一。

答案 1 :(得分:1)

我遇到了同样的问题,但我的tf.model_variables()是空的。除了用tf.trainable_variables()替换tf.model_variables()之外,我只能使用相同的想法保存必要的变量。这很好地排除了ADAM变量。我的模型保存代码如下所示:

saver = tf.train.Saver(tf.trainable_variables())
saver.save(session, 'model.ckpt')

答案 2 :(得分:0)

您可以在保护程序中指定变量:

saver = tf.train.Saver(...variables...) saver.save(sess, 'my-model', global_step=step)

参考:https://www.tensorflow.org/api_docs/python/tf/train/Saver