在学习机器学习的时候,很多时候我们并不要求计算机有很高的单核算力,反而要求计算机有比较高的并行计算功能。这时候我们就可以使用GPU(显卡)来帮我们运行机器学习的代码了。业内比较出名的机器学习的框架有TensorFlow和Keras,那么TensorFlow使用GPU和Keras使用GPU需要设置些什么呢?阅读这篇文章你会得到了解。
1. 训练运行时候指定GPU
运行时候加一行代码:
CUDA_VISIBLE_DEVICES=1 python train.py
2. 运行过程中按需或者定量分配GPU
tensorflow直接在开启Session时候加几行代码就行,而Keras指定GPU,并限制按需用量和TensorFlow不太一样,因为keras训练是封装好的,不好对Session操作。如下是两种对应的操作。
keras中的操作:
import os import tensorflow as tf from keras.backend.tensorflow_backend import set_session # 指定第一块GPU可用 os.environ["CUDA_VISIBLE_DEVICES"] = "0" #指定GPU的第二种方法 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 #定量 config.gpu_options.allow_growth = True #按需 set_session(tf.Session(config=config))
TensorFlow中的操作:
#指定GPU import os os.environ["CUDA_VISIBLE_DEVICES"] = "0" #设置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)
补充:Keras以及Tensorflow强制使用CPU,GPU
Keras如果是使用Theano后端的话,应该是自动不使用GPU只是用CPU的,启动GPU使用Theano内部命令即可。
对于Tensorflow后端的Keras以及Tensorflow会自动使用可见的GPU,而我需要其必须只运行在CPU上。网上查到三种方法,最后一种方法对我有用,但也对三种都做如下记录:
使用tensorflow的 with tf.device('/cpu:0'):函数。简单操作就是把所有命令都放在前面所述的域里面。
使用tensorflow声明Session时的参数: 关于tensorflow中Session中的部分参数设置,以及Keras如何设置其调用的Tensorflow的Session,可以参见Keras设定GPU使用内存大小(Tensorflow backend)。
对于Tensorflow,声明Session的时候加入device_count={'gpu':0}即可,代码如下:
import tensorflow as tf sess = tf.Session(config=tf.ConfigProto(device_count={'gpu':0}))
对于Keras,则调用后端函数,设置其使用如上定义的Session即可,代码如下:
import tensorflow as tf import keras.backend.tensorflow_backend as KTF KTF.set_session(tf.Session(config=tf.ConfigProto(device_count={'gpu':0})))
对于多线程以及GPU内存设置等可以参见Keras设定GPU使用内存大小(Tensorflow backend)。
3、第三种是使用CUDA_VISIBLE_DEVICES命令行参数,代码如下:
CUDA_VISIBLE_DEVICES="0" python3 train.py
以上就是TensorFlow使用GPU和keras使用GPU时需要注意的细节了,希望能给大家一个参考,也希望大家多多支持W3Cschool。