在學(xué)習(xí)機(jī)器學(xué)習(xí)的時(shí)候,很多時(shí)候我們并不要求計(jì)算機(jī)有很高的單核算力,反而要求計(jì)算機(jī)有比較高的并行計(jì)算功能。這時(shí)候我們就可以使用GPU(顯卡)來幫我們運(yùn)行機(jī)器學(xué)習(xí)的代碼了。業(yè)內(nèi)比較出名的機(jī)器學(xué)習(xí)的框架有TensorFlow和Keras,那么TensorFlow使用GPU和Keras使用GPU需要設(shè)置些什么呢?閱讀這篇文章你會(huì)得到了解。
1. 訓(xùn)練運(yùn)行時(shí)候指定GPU
運(yùn)行時(shí)候加一行代碼:
CUDA_VISIBLE_DEVICES=1 python train.py
2. 運(yùn)行過程中按需或者定量分配GPU
tensorflow直接在開啟Session時(shí)候加幾行代碼就行,而Keras指定GPU,并限制按需用量和TensorFlow不太一樣,因?yàn)閗eras訓(xùn)練是封裝好的,不好對(duì)Session操作。如下是兩種對(duì)應(yīng)的操作。
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" #設(shè)置GPU定量分配 config = tf.ConfigProto() config.gpu_options.per_process_gpu_memory_fraction = 0.9 # 占用GPU90%的顯存 session = tf.Session(config=config) #設(shè)置GPU按需分配 config = tf.ConfigProto() config.gpu_options.allow_growth = True session = tf.Session(config=config)
補(bǔ)充:Keras以及Tensorflow強(qiáng)制使用CPU,GPU
Keras如果是使用Theano后端的話,應(yīng)該是自動(dòng)不使用GPU只是用CPU的,啟動(dòng)GPU使用Theano內(nèi)部命令即可。
對(duì)于Tensorflow后端的Keras以及Tensorflow會(huì)自動(dòng)使用可見的GPU,而我需要其必須只運(yùn)行在CPU上。網(wǎng)上查到三種方法,最后一種方法對(duì)我有用,但也對(duì)三種都做如下記錄:
使用tensorflow的 with tf.device('/cpu:0'):函數(shù)。簡(jiǎn)單操作就是把所有命令都放在前面所述的域里面。
使用tensorflow聲明Session時(shí)的參數(shù): 關(guān)于tensorflow中Session中的部分參數(shù)設(shè)置,以及Keras如何設(shè)置其調(diào)用的Tensorflow的Session,可以參見Keras設(shè)定GPU使用內(nèi)存大小(Tensorflow backend)。
對(duì)于Tensorflow,聲明Session的時(shí)候加入device_count={'gpu':0}即可,代碼如下:
import tensorflow as tf sess = tf.Session(config=tf.ConfigProto(device_count={'gpu':0}))
對(duì)于Keras,則調(diào)用后端函數(shù),設(shè)置其使用如上定義的Session即可,代碼如下:
import tensorflow as tf import keras.backend.tensorflow_backend as KTF KTF.set_session(tf.Session(config=tf.ConfigProto(device_count={'gpu':0})))
對(duì)于多線程以及GPU內(nèi)存設(shè)置等可以參見Keras設(shè)定GPU使用內(nèi)存大小(Tensorflow backend)。
3、第三種是使用CUDA_VISIBLE_DEVICES命令行參數(shù),代碼如下:
CUDA_VISIBLE_DEVICES="0" python3 train.py
以上就是TensorFlow使用GPU和keras使用GPU時(shí)需要注意的細(xì)節(jié)了,希望能給大家一個(gè)參考,也希望大家多多支持W3Cschool。