TensorFlow中使用GPU

TensorFlow默認會佔用設備上全部的GPU以及每一個GPU的全部顯存；若是指定了某塊GPU,也會默認一次性佔用該GPU的全部顯存。能夠經過如下方式解決：

1 Python代碼中設置環境變量，指定GPU

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(...)