初步尝试 Keras (基于 Tensorflow 后端)深度框架时, 发现其对于 GPU 的使用比较神奇, 默认竟然是全部占满显存, 1080Ti 跑个小分类问题, 就一下子满了. 而且是服务器上的两张 1080Ti.

服务器上的多张 GPU 都占满, 有点浪费性能.

因此, 需要类似于 Caffe 等框架的可以设定 GPU ID 和显存自动按需分配.

实际中发现, Keras 还可以限制 GPU 显存占用量.

这里涉及到的内容有:

• GPU ID 设定
• GPU 显存占用按需分配
• GPU 显存占用限制
• GPU 显存优化

1. GPU ID 设定

#! -- coding: utf-8 --*--
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

这里将 GPU ID 设为 1.

GPU ID 从 0 开始, GPUID=1 即表示第二块 GPU.

2. GPU 显存占用按需分配

#! -- coding: utf-8 --*--
import tensorflow as tf
import keras.backend.tensorflow_backend as ktf

# GPU 显存自动调用
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
session = tf.Session(config=config)
ktf.set_session(session)

3. GPU 显存占用限制

#! -- coding: utf-8 --*--
import tensorflow as tf
import keras.backend.tensorflow_backend as ktf

# 设定 GPU 显存占用比例为 0.3
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.3
session = tf.Session(config=config)
ktf.set_session(session )

这里虽然是设定了 GPU 显存占用的限制比例(0.3), 但如果训练所需实际显存占用超过该比例, 仍能正常训练, 类似于了按需分配.

设定 GPU 显存占用比例实际上是避免一定的显存资源浪费.

4. GPU ID 设定与显存按需分配

#! -- coding: utf-8 --*--
import os
import tensorflow as tf
import keras.backend.tensorflow_backend as ktf

# GPU 显存自动分配
config = tf.ConfigProto()
config.gpu_options.allow_growth=True
#config.gpu_options.per_process_gpu_memory_fraction = 0.3
session = tf.Session(config=config)
ktf.set_session(session)

# 指定GPUID, 第一块GPU可用
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

5. 利用fit_generator最小化显存占用比例/数据Batch化

#! -- coding: utf-8 --*--

# 将内存中的数据分批(batch_size)送到显存中进行运算
def generate_arrays_from_memory(data_train, labels_train, batch_size):
    x = data_train
    y=labels_train
    ylen=len(y)
    loopcount=ylen // batch_size
    while True:
        i = np.random.randint(0,loopcount)
        yield x[i*batch_size:(i+1)*batch_size],y[i*batch_size:(i+1)*batch_size]

# load数据到内存
data_train=np.loadtxt("./data_train.txt")
labels_train=np.loadtxt('./labels_train.txt')
data_val=np.loadtxt('./data_val.txt')
labels_val=np.loadtxt('./labels_val.txt')

hist=model.fit_generator(generate_arrays_from_memory(data_train,
                                                     labels_train,
                                                     batch_size),
                         steps_per_epoch=int(train_size/bs),
                         epochs=ne,
                         validation_data=(data_val,labels_val),
                         callbacks=callbacks )

5.1 数据 Batch 化

#! -- coding: utf-8 --*--

def process_line(line):  
    tmp = [int(val) for val in line.strip().split(',')]  
    x = np.array(tmp[:-1])  
    y = np.array(tmp[-1:])  
    return x,y  

def generate_arrays_from_file(path,batch_size):  
    while 1:  
        f = open(path)  
        cnt = 0  
        X =[]  
        Y =[]  
        for line in f:  
            # create Numpy arrays of input data  
            # and labels, from each line in the file  
            x, y = process_line(line)  
            X.append(x)  
            Y.append(y)  
            cnt += 1  
            if cnt==batch_size:  
                cnt = 0  
                yield (np.array(X), np.array(Y))  
                X = []  
                Y = []  
    f.close()  

Reference

[1] - keras深度训练4:GPU设置