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
以上為個人經驗,希望能給大家一個參考,也希望大家多多支持腳本之家。
您可能感興趣的文章:- TensorFlow2.0使用keras訓練模型的實現
- tensorflow2.0教程之Keras快速入門
- 解決TensorFlow調用Keras庫函數存在的問題
- Keras模型轉成tensorflow的.pb操作
- 完美解決TensorFlow和Keras大數據量內存溢出的問題
- keras和tensorflow使用fit_generator 批次訓練操作
- 解決tensorflow 與keras 混用之坑