使用Tensorflow时将大矩阵计算拟合到内存中

Question

我正在尝试构建一个有两个阶段的模型。 第一个采用输入图像并将其传递给conv-deconv网络。得到的Tensor具有对应于期望输出图像中的像素的条目（与输入图像大小相同）。

要计算最终输出图像，我想从第一阶段获取在每个像素位置生成的值，并将其用作应用于整个输入图像的缩小函数的附加输入。第二步没有可训练的变量，但它的计算/内存成本会随着输入的大小呈指数增长（每个输出像素都是所有输入像素的函数）。

我目前正在使用tf.map_fn来计算输出图像。我将输出像素计算功能映射到第一阶段的结果。我的愿望是张量流将分配存储器来存储每个像素计算所需的中间张量，然后在继续下一个像素计算之前释放该存储器。但相反，似乎永远不会释放导致OOM错误的中间计算。

有没有告诉tensorflow（显式或隐式）它应该释放分配的内存以保存计算中不再需要的Tensor数据？

Answer 1

只要未来的计算不再需要张量，TensorFlow就会为张量释放内存。您可以通过查看内存释放消息来验证这一点，如notebook。

所示

由于TensorFlow以低效的顺序执行内存中的节点，因此可能内存不足。

例如，请考虑以下计算：

k = 2000
a = tf.random_uniform(shape=(k,k))
for i in range(n):
    a = tf.matmul(a, tf.random_uniform(shape=(k,k)))

评估的顺序如下所示

首先评估所有圈子（tf.random_uniform）节点，然后是方块（tf.matmul）。对于最佳订单，这与O(n)相比具有O(1)内存要求。

您可以使用控制依赖项强制执行特定的执行顺序，即使用如下的帮助函数：

import tensorflow.contrib.graph_editor as ge
def run_after(a_tensor, b_tensor):
    """Force a to run after b"""
    ge.reroute.add_control_inputs(a_tensor.op, [b_tensor.op])