通过Keras / TF 2019限制GPU内存使用吗?

时间:2019-04-22 03:45:31

标签: tensorflow keras deep-learning gpu

我已阅读以下答案:

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.2
set_session(tf.Session(config=config))

但是它根本不起作用。 keras和TF中似乎都有太多更新,以至于2017年编写的几乎所有内容都不起作用! 那么,如何限制内存使用量呢?

2 个答案:

答案 0 :(得分:1)

似乎我有同样的问题。你尝试过吗?

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
set_session(sess)

此方法将使应用程序根据运行时分配仅分配尽可能多的GPU内存。

答案 1 :(得分:1)

一种限制在张量流中保留所有GPU RAM的方法是增加保留量。这种方法将允许您使用相同的GPU训练多个NN,但无法为要保留的内存量设置阈值。

在导入keras之前使用以下代码段,或仅使用tf.keras

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)

    except RuntimeError as e:
        print(e)