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

Question

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

Answer 1

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

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

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

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