如何在 TF 2.1 上设置动态内存增长？

Posted 2023-03-16

【中文标题】如何在 TF 2.1 上设置动态内存增长？

【英文标题】：How to set dynamic memory growth on TF 2.1?

【发布时间】：2020-05-19 17:51:37

【问题描述】：

使用以前版本的 tensorflow+keras，我可以设置一个“allow_growth”选项并使用 nvidia-smi 查看实时内存使用情况。否则，它将立即被进程分配。现在，在 tensorflow 2.1 中使用 tf.keras 我找不到这样做的方法。任何帮助表示赞赏！

【问题讨论】：

这能回答你的问题吗？ tensorflow.org/api_docs/python/tf/config/experimental/…

我会试试这个并报告，谢谢！

是的，这行得通。谢谢！

【参考方案1】：

如果您有多个 GPU，您将只允许第一个 GPU 的内存增长。

physical_devices = tf.config.list_physical_devices('GPU') 
tf.config.experimental.set_memory_growth(physical_devices[0], True)

如果您想为所有 GPU 执行此操作，则需要为每个实例设置它。

physical_devices = tf.config.list_physical_devices('GPU') 
for gpu_instance in physical_devices: 
    tf.config.experimental.set_memory_growth(gpu_instance, True)

【讨论】：

是否有不依赖于框架的驱动级解决方案？基本上说 - 默认情况下阻止所有 CUDA 应用程序占用所有 GPU 内存？

您可以通过设置 CUDA_VISIBLE_DEVICES 来获得此行为。我没有用其他应用程序/框架对此进行测试，但我相信它应该可以工作。 developer.nvidia.com/blog/…

不，这只是从代码中完全隐藏了 GPU。我希望代码不要在 GPU 上贪婪地分配内存，仅此而已。

哦，我明白了，不幸的是，我怀疑这样的事情是否存在。

我看到即使没有设置，TF2 也不会占用所有内存。所以也许这不需要？