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

Question

在推理过程中，我们不需要在通过网络传播时保留前一层的激活。但是，由于我们没有明确告诉程序丢弃它们，因此它不会区分训练和推理过程。有没有一种方法——也许是一个简单的标志、类、方法——在 Tensorflow 中进行这种内存管理？简单地使用 tf.stop_gradient 行吗？

Answer 1

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

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

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

一旦消除了向后传递，tensorflow 就可以优化其内存使用，尤其是自动优化 free or reuse memory taken by unused nodes .