tensorflow指定GPU - xd21303/Notebook GitHub Wiki

在.bashrc中加入

export CUDA_VISIBLE_DEVICES=0

0是第一个显卡,1是第二块显卡,以此类推。

或者在代码前加入

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

不使用GPU,使用CPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

tensorflow

为显存分配使用比例

在建立tf.Session加入设置数据(显存使用比例为1/3),但有时你虽然设置了使用上限,在程序需要更高显存时还是会越过该限制

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)  
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))  

自适应分配

会自适应分配显存,不会将显存全部分配导致资源浪费

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

keras

为显存分配使用比例

import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.333
session = tf.Session(config=config)
KTF.set_session(session)

自适应分配

import keras.backend.tensorflow_backend as KTF

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

如有设置fit_generator

将多线程关闭

可将

use_multiprocessing=True

改为

use_multiprocessing=False