tensorflow - tensorflow 消耗的 GPU 内存是否完全等于所需的

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

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

标签： tensorflowtensorflow-c++

不幸的是，没有开箱即用的标志，但这可以（手动）完成：

默认情况下，TF 会分配所有可用的 GPU 内存。设置set_allow_growth为true，导致 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 的值将实现您想要的。