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

Question

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

machine-learningtensorflowgpudeep-learning

9

有没有一种简单的方法可以找到由tensorflow初始化的inception-resnet-v2模型所消耗的GPU内存？这包括所需的推理和反向传播内存。

- kwotsin

2个回答

2

由于目前使用 gpu.options.allow_growth 和 gpu_options.per_process_gpu_memory_fraction 进行模型大小估计仍然是一种试错和繁琐的解决方案，我建议结合 tensorboard 使用 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)

运行您的模型和TensorBoard，导航到所需部分的图形并阅读节点统计信息。

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

- whiletrue

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Dmytro Prylipko · Accepted Answer

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