在學習機器學習的時候，很多時候我們并不要求計算機有很高的單核算力，反而要求計算機有比較高的并行計算功能。這時候我們就可以使用GPU（顯卡）來幫我們運行機器學習的代碼了。業(yè)內比較出名的機器學習的框架有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后端的話，應該是自動不使用ＧＰＵ只是用ＣＰＵ的，啟動ＧＰＵ使用Theano內部命令即可。
對于Tensorflow后端的Keras以及Tensorflow會自動使用可見的ＧＰＵ，而我需要其必須只運行在ＣＰＵ上。網(wǎng)上查到三種方法，最后一種方法對我有用，但也對三種都做如下記錄：
使用tensorflow的 with tf.device('/cpu:0'):函數(shù)。簡單操作就是把所有命令都放在前面所述的域里面。
使用tensorflow聲明Session時的參數(shù)： 關于tensorflow中Session中的部分參數(shù)設置，以及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,則調用后端函數(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})))
對于多線程以及ＧＰＵ內存設置等可以參見Keras設定GPU使用內存大小(Tensorflow backend)。
3、第三種是使用CUDA_VISIBLE_DEVICES命令行參數(shù)，代碼如下：
CUDA_VISIBLE_DEVICES="0" python3 train.py
以上就是TensorFlow使用GPU和keras使用GPU時需要注意的細節(jié)了，希望能給大家一個參考，也希望大家多多支持W3Cschool。
