Tensorflow内存管理分块？

Question

我在计算一个网络函数的约化时遇到了一个问题，这个网络由一个大的（20000x200000）矩阵表示，这个矩阵生成为一对点之间的距离矩阵。在

最小示例，输入一个200000x2笛卡尔坐标数组：

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：

^{pr2}$

因此，只有行和存储为张量，而不是完整的距离矩阵。然而，尽管这种方法在CPU上工作得很好，但在GPU上却陷入停顿。在

Answer 1

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

通过使用变量和assign_add递增地更新内容而不需要创建许多中间张量，可以提高内存效率。有一个计算“所有成对距离”here的替代公式，它可能更容易转换成这种形式