Tensorflow：并行for循环导致内存不足

Question

我的张量流代码如下：

for i in range(100):
    Zk = Zk + function_call(...)

似乎tensorflow并行运行这100次迭代，并保留许多与Zk具有相同大小的临时向量。但是，由于Zk是一个非常长的向量，因此会立即导致内存不足错误。

任何人都可以就如何强制Tensorflow顺序执行循环提出一些建议。非常感谢。

Answer 1

解决此问题的最简单方法是使用with tf.control_dependencies():块：

Zk = ...
for i in range(100):
    with tf.control_dependencies([Zk.op]):
        Zk = Zk + function_call(...)

with tf.control_dependencies():块确保function_call()创建的操作不会运行，直到计算出Zk的上一个值。这有效地导致循环迭代按顺序运行。

另一种解决方案涉及使用tf.while_loop()来定义迭代，而不是Python循环。 tf.while_loop()函数有一个parallel_iterations可选参数，允许您减少循环的独立迭代之间的并行度。这种方法的一个优点是它可以使图形保持较小：循环将使用O（1）节点而不是O（N）来执行N次迭代。但是，将任意Python function_call()重新转换为TensorFlow循环体可能有点棘手，因此您可能会发现最初使用with tf.control_dependencies():更容易。