内存效率的设计技巧TensorFlow NN

Question

我在CPU上使用TensorFlow运行一个简单的卷积神经网络，使用Python / Numpy从磁盘加载数据并准备它。 Python / Numpy部分消耗很少的内存，正如预期的那样，因为我只加载了32个1MB的图像。

然而，在天文数据下运行sess.run(AdamOptimizer)内存使用峰值时，可以在下面的memory_profiler输出中看到，显示4个训练周期和一个测试周期。

12GB峰值限制了我的批次大小，同时训练网络的数量或我的图层的大小。我认为这个峰值与从Python到TF或正向激活的数据传输有关，因为它在sess.run(accuracy)语句中等效存在。

我是否设计了CNN图表？有人可以指出我应该记住的设计规则，特别是对于记忆吗？有可能及时涂抹这个吗？ - 如果12GB峰值的时间长一倍，高一半则会有所改善。最后有没有办法发出峰值开始和结束时间的信号，这样我就可以管理多个网络来训练隔行扫描？ （请注意，狭窄的2秒峰值出现在单个sess.run()语句中，大约需要20秒）

Answer 1

如果有人要添加任何内容，我很乐意转移答案。只是想分享我找到的东西。

似乎tf.nn.conv2d特别是记忆力饥渴，也许它在某种程度上是完全有道理的，但对我而言，它远远超出预期。以下简单程序使用5.5GB，我不希望超过几百MB：

x = tf.placeholder(tf.float32, shape=[None, 256, 256, 4])
w = tf.Variable(tf.truncated_normal([f, f, 4, 8], stddev=0.1))
outp = tf.nn.conv2d(x, w, strides=[1,1,1,1], padding='SAME')

outp_v = sess.run(outp,
       feed_dict={x:np.ones(32*256*256*4).reshape(32,256,256,4)})

与内存分析器一样： 对于此输出，我使用了相当大的数字f = 16，但我发现内存使用与它无关。除了足够小f之外，计算速度太快以至于它低于内存配置文件实用程序的时间分辨率，并且只有一个尖峰（没有平台）保持在任意高度。

其他人发现了类似的内存使用情况： tensorflow conv2d memory consumption explain?

Answer 2

您选择的优化程序是内存高峰达到4倍的原因。亚当每个卷积权重使用8位。我打算按照Google的建议使用其他优化器，例如AdaFactorOptimizer

参考： 子标题：大量模型权重 https://cloud.google.com/tpu/docs/troubleshooting