tensorflow - 如何在 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)