TensorFlow中使用GPU
阿新 • • 發佈:2020-02-16
TensorFlow預設會佔用裝置上所有的GPU以及每個GPU的所有視訊記憶體;如果指定了某塊GPU,也會預設一次性佔用該GPU的所有視訊記憶體。可以通過以下方式解決:
1 Python程式碼中設定環境變數,指定GPU
本文所有程式碼在tensorflow 1.12.0中測試通過。
import os os.environ["CUDA_VISIBLE_DEVICES"] = "2" # 指定只是用第三塊GPU
2 系統環境變數中指定GPU
# 只使用第2塊GPU,在demo_code.py,機器上的第二塊GPU變成”/gpu:0“,不過在執行時所有的/gpu:0的運算將被放到第二塊GPU上 CUDA_VISIBLE_DEVICES=1 python demo_code.py #只使用第一塊GPU和第二塊GPU CUDA_VISIBLE_DEVICES=0,1 python demo_code.py
3 動態分配GPU視訊記憶體
# allow_soft_placement=True 沒有GPU的話在CPU上執行 config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True) config.gpu_options.allow_growth = True # 按需分配視訊記憶體 with tf.Session(config=config) as sess: sess.run(...)
4 按固定比例分配視訊記憶體
# 按照固定的比例分配。 config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True) # 以下程式碼會佔用所有可使用的GPU的40%視訊記憶體 config.gpu_options.per_process_gpu_memory_fraction = 0.4 with tf.Session(config=config) as sess: sess.run(...)
在我的裝置中設定後GPU佔用情況如下:
gz_6237_gpu Sat Feb 15 23:01:56 2020 418.87.00 [0] GeForce RTX 2080 Ti | 43'C, 0 % | 4691 / 10989 MB | dc:python/1641(4681M)
5 通過tf.device將運算指定到特定裝置上
with tf.device("/gpu:0"): b = tf.Variable(tf.zeros([1])) W = tf.Variable(tf.random_uniform([1, 2], -1.0, 1.0)) y = tf.matmul(W, x_data) + b
這種方式不推薦。TF的kernel中國定義了哪些操作可以跑在GPU上,哪些不可以,因此強制指定GPU會降低程式的可移植性。
推薦的做法是:在建立會話時,指定引數allow_soft_placement=True;這樣如果運算無法在GPU上執行,TF會自動將它放在CPU上執行。
config = tf.ConfigProto(allow_soft_placement=True) with tf.Session(config=config) as sess: sess.run(...)