1

我使用 tensorflow c++ 版本进行 CNN 推理。我已经set_allow_growth(true),但它仍然消耗比实际需要更多的 GPU 内存。

set_per_process_gpu_memory_fraction只能设置GPU显存的上限,但不同的CNN模型有不同的上限。有没有解决问题的好方法

4

1 回答 1

1

不幸的是,没有开箱即用的标志,但这可以(手动)完成:

默认情况下,TF 会分配所有可用的 GPU 内存。设置set_allow_growthtrue,导致 TF 以块的形式分配所需的内存,而不是一次分配所有 GPU 内存。每次 TF 需要比已分配更多的 GPU 内存时,它都会分配另一个块。

此外,正如您所提到的,TF 支持set_per_process_gpu_memory_fraction以占总 GPU 内存的百分比来指定进程可能需要的最大 GPU 内存。这会导致内存不足 (OOM) 异常,以防 TF 需要比允许更多的 GPU 内存。

不幸的是,我认为块大小不能由用户设置,并且在 TF 中是硬编码的(出于某种原因,我认为块大小是 4GB,但我不确定)。

这导致能够指定允许 TF 使用的最大 GPU 内存量(以百分比表示)。如果您知道您总共有多少 GPU 内存(可以通过 检索nvidia-smi,并且您知道您想要允许多少内存,您可以用百分比计算它并将其设置为 TF.

如果您运行少量的神经网络,您可以通过使用不同的允许 GPU 内存运行它来找到每个神经网络所需的 GPU 内存,例如二分搜索,并查看使 NN 能够运行的最小分数是多少。然后,将您找到的值设置set_per_process_gpu_memory_fraction为每个 NN 的值将实现您想要的。

于 2020-01-24T09:28:32.510 回答