Keras和TensorFlow设置GPU及其使用率

Keras和TensorFlow设置GPU及其使用率

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

1、TensorFlow1.X

1.1指定GPU

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

tf.device('/gpu:1')

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

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

通过以上代码，只使得编号为1的GPU对当前python程序是可见的。
如果要指定多块GPU，则将相应GPU编号用逗号隔开即可：

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 tf
from keras.backend.tensorflow_backend import set_session
config = 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.3
set_session(tf.Session(config=config)) 

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

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
config = 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 = True
set_session(tf.Session(config=config)) 

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

指定占GPU内存比例：

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session
gpus = 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 tf
from keras.backend.tensorflow_backend import set_session
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)