Tensorflow 内存管理——分块？

Question

我在计算网络函数的约简时遇到问题，该网络由一个大 (200000x200000) 矩阵表示，该矩阵生成为点对之间的距离矩阵。

最小的例子，输入 X 一个 200000x2 的笛卡尔坐标的 numpy 数组：

x = tf.constant(X[:,0], shape=[X.shape[0],1])
y = tf.constant(X[:,1], shape=[X.shape[0],1])
dx = x - tf.transpose(x)
dy = y - tf.transpose(y)
D = tf.sqrt(dx*dx + dy*dy)
M = 0.1 * 5.0 / tf.pow(4.0 + D, 1.5)
res = tf.reduce_sum(betaM)

在 CPU 上运行时，内存（我的 MBP 上为 16GB）很快就被超额认购了，系统陷入了停顿。据推测，tf 正试图将整个 D（和 M？）存储在内存中。

如果我在 C/C++ 中编写此代码，我很可能会遍历矩阵行，在执行过程中对每一行求和，并且从不存储整个矩阵。 GPU 也是如此——我将细分（虚拟）矩阵并以块的形式执行缩减。

有没有什么技巧可以让 tf 遵循更加分块的行为，从而节省内存？

干杯，

克里斯

编辑：

解决内存问题的另一种方法是使用tf.map_fn：

rowsums = tf.map_fn(lambda i: tf.reduce_sum(tf.sqrt(tf.reduce_sum(tf.pow(i - x,2),1))) , x)
res = tf.reduce_sum(rowsums)

因此，只有行和存储为张量，而不是全距离矩阵。然而，尽管这种方法在 CPU 上运行良好，但在 GPU 上却停滞不前。

Answer 1

这里真正需要（但尚未实现）的是 cwise 融合。现在发生的事情是2*sqrt(a+b) 将为a+b 分配新张量，然后为sqrt 分配新张量，然后为2*sqrt 分配另一个张量。 PS，你可以通过检查内存分配messages（需要verbose logging）来挖掘内存的去向

您可以通过使用变量和assign_add 来提高内存效率，从而在不创建许多中间张量的情况下逐步更新内容。有一个计算“所有成对距离”here 的替代公式可能更容易转换为这种形式