Tensorflow:如何减少仅用于推理模型的内存占用？

Question

在推理过程中，当我们在网络中传播时，我们不需要保留前几层的激活。然而，由于我们没有显式地告诉程序丢弃它们，所以它不区分训练和推理过程。有没有一种方法-perhaps一个简单的标志，类，方法-在Tensorflow中做这种内存管理？简单地使用tf.stop_gradient就行了吗？

Answer 1

最简单的方法是使用freeze_graph.py脚本"freeze“(tensorflow的术语)您的模型。

这个脚本基本上删除了所有不必要的操作，还用常量替换了所有变量，然后导出回磁盘上的结果图。

为此，您需要在图形中指定哪些是您在推理过程中使用的输出。无法到达输出的节点(可能是汇总、损失、梯度等)会被自动丢弃。

一旦消除了反向遍历，tensorflow可以优化其内存使用，特别是自动执行free or reuse memory taken by unused nodes。