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

Question

对于以前版本的 tensorflow+keras，我能够设置一个 'allow_growth' 选项并使用 nvidia-smi 查看实时内存使用情况。否则它将全部被进程立即分配。现在，在 tensorflow 2.1 中使用 tf.keras 我找不到办法做到这一点。感谢您的帮助！

Answer 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)

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

How to set dynamic memory growth on TF 2.1?

memory

gpu

tensorflow

tf.keras