如何计算在TensorFlow中运行模型所需的GPU内存？

Question

是否有一种直接的方法来查找由张量流初始化的inception-resnet-v2模型消耗的GPU内存？这包括推理和所需的后备存储器.

Answer 1

您可以明确计算存储参数所需的内存，但恐怕很难计算出训练所需的所有缓冲区的大小。也许，更聪明的方法是让 TF 为你做这件事。将gpu_options.allow_growth配置选项设置为 True 并查看它消耗了多少。另一种选择是尝试较小的值，gpu_options.per_process_gpu_memory_fraction直到它因内存不足而失败。

Answer 2

由于使用gpu.options.allow_growth和gpu_options.per_process_gpu_memory_fraction进行模型大小估计目前是一个反复试验且乏味的解决方案，因此我建议tf.RunMetadata()与张量板结合使用。

例子：

run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
summary, _ = sess.run(train_step, feed_dict, options=run_options, run_metadata=run_metadata)

train_writer.add_run_metadata(run_metadata, 'step%d' % i)

运行模型和张量板，导航到图形的所需部分并读取节点统计信息。

来源： https: //www.tensorflow.org/get_started/graph_viz