本文共 2277 字，大约阅读时间需要 7 分钟。

Keras和TensorFlow设置GPU及其使用率

本文讲一下深度学习框架Keras和Tensorflow如何指定GPU训练，及如何设置训练任务占用的GPU内存。

1、TensorFlow1.X

1.1指定GPU

直接在代码中用以下语句指定并不能实现我们的目的：

tf.device('/gpu:1')

需要在python代码中通过CUDA_VISIBLE_DEVICES来指定：

import osos.environ['CUDA_VISIBLE_DEVICES']='1'

通过以上代码，只使得编号为1的GPU对当前python程序是可见的。

os.environ['CUDA_VISIBLE_DEVICES']='2,3'

也可以在终端中，来指定想使用的GPU，将指定GPU的代码放在python程序的前面：

CUDA_VISIBLE_DEVICES=1 python test.py

1.2 指定GPU使用量

config = tf.ConfigProto() config.gpu_options.per_process_gpu_memory_fraction = 0.9 # 占用GPU90%的显存 session = tf.Session(config=config)

设置GPU使用量最小。

config = tf.ConfigProto() config.gpu_options.allow_growth = True session = tf.Session(config=config)

2、TensorFlow2.X

2.1指定GPU

与TF1.X方法基本一样

2.2 指定GPU使用量

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')for gpu in gpus:    tf.config.experimental.per_process_gpu_memory_fraction = 0.9

设置GPU使用量最小。

gpus = tf.config.experimental.list_physical_devices(device_type='GPU')for gpu in gpus:    tf.config.experimental.set_memory_growth(gpu, True)

3、Keras

3.1 指定GPU

方法与TensorFlow相同。

3.2 指定GPU使用量（TF1.X）

指定占GPU内存比例：

import tensorflow as tffrom keras.backend.tensorflow_backend import set_sessionconfig = tf.ConfigProto()config.gpu_options.allocator_type = 'BFC' #A "Best-fit with coalescing" algorithm, simplified from a version of dlmalloc.config.gpu_options.per_process_gpu_memory_fraction = 0.3set_session(tf.Session(config=config))

根据实际使用情况分配最小内存：

import tensorflow as tffrom keras.backend.tensorflow_backend import set_sessionconfig = tf.ConfigProto()config.gpu_options.allocator_type = 'BFC' #A "Best-fit with coalescing" algorithm, simplified from a version of dlmalloc.config.gpu_options.allow_growth = Trueset_session(tf.Session(config=config))

3.3 指定GPU使用量（TF2.X）

指定占GPU内存比例：

import tensorflow as tffrom keras.backend.tensorflow_backend import set_sessiongpus = tf.config.experimental.list_physical_devices(device_type='GPU')for gpu in gpus:    tf.config.experimental.per_process_gpu_memory_fraction = 0.9

根据实际使用情况分配最小内存：

import tensorflow as tffrom keras.backend.tensorflow_backend import set_sessiongpus = tf.config.experimental.list_physical_devices(device_type='GPU')for gpu in gpus:    tf.config.experimental.set_memory_growth(gpu, True)

转载地址：http://zvrti.baihongyu.com/