如何在 TF 2.1 上设置动态内存增长? [英] How to set dynamic memory growth on TF 2.1?

问题描述

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

With previous versions of tensorflow+keras I was able to set an 'allow_growth' option and view realtime memory usage with nvidia-smi. Otherwise it will all gett allocated immediately by the process. Now, using tf.keras in tensorflow 2.1 I can't find a way to do this. Any help is appreciated!

推荐答案

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

In case you have several GPUs, you will allow memory growth only for the first GPU.

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

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

If you want to do it for all GPUs you need to set it for every instance.

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 上设置动态内存增长?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！