本篇介紹如何指定 TensorFlow 與 Keras 程式所使用的 GPU 顯示卡與記憶體用量。

在 TensorFlow 或 Keras 中使用 NVIDIA 的 GPU 做運算時,預設會把整台機器上所有的 GPU 卡都獨佔下來,而且不管實際需要多少顯示卡的記憶體,每張卡的記憶體都會被佔滿,以下介紹如何調整設定,讓多張顯示卡可以分給多個程式或多人使用。

指定 GPU 顯示卡

若要只使用特定的 GPU 卡,可以使用 CUDA_VISIBLE_DEVICES 這個環境變數來設定,它的值代表 CUDA 程式可以使用的 GPU 卡編號(從 0 開始),例如只讓 CUDA 程式使用第一張 GPU 卡:

# 只讓 CUDA 程式使用第一張 GPU 卡
export CUDA_VISIBLE_DEVICES=0
python my_script.py

這樣當 my_script.py 這個 Python 程式在執行時,就只會用到機器上的第一張 GPU 卡。若要指定多張 GPU 卡,則以逗號分隔:

# 讓 CUDA 程式使用第一張與第三張 GPU 卡
export CUDA_VISIBLE_DEVICES=0,2
python my_script.py

如果自己會用的 GPU 卡都是固定的,我們可以將 CUDA_VISIBLE_DEVICES 的設定寫在 ~/.bashrc 中,在登入 Linux 系統時就自動設定好。而如果要讓不同的程式用不同的 GPU 卡計算,分散計算量的話,可以在執行程式時直接以 CUDA_VISIBLE_DEVICES 指定:

# 使用第一張 GPU 卡
CUDA_VISIBLE_DEVICES=0 python my_script1.py

# 使用第二張與第三張 GPU 卡
CUDA_VISIBLE_DEVICES=1,2 python my_script2.py

另外還有一種方式是直接在 Python 程式中更改 CUDA_VISIBLE_DEVICES 這個環境變數,此種方式的原理也是一樣的,只是這樣可以把 GPU 卡的指定邏輯寫在 Python 程式中:

import os

# 使用第一張與第三張 GPU 卡
os.environ["CUDA_VISIBLE_DEVICES"] = "0,2"

指定 GPU 顯示卡記憶體用量上限

若在 TensorFlow 中,我們可以使用 tf.GPUOptions 來調整程式佔用的 GPU 記憶體:

import tensorflow as tf
W = tf.constant([1.0, 2.0, 3.0, 4.0], shape=[2, 2], name='W')
x = tf.constant([1.3, 2.4], shape=[2, 1], name='x')
y = tf.matmul(W, x)

# 只使用 30% 的 GPU 記憶體
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.3)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
print(sess.run(y))

在以 TensorFlow 為 backend 的 Keras 程式中,我們可以透過以下的設定方式來指定 GPU 記憶體的佔用量:

import tensorflow as tf

# 只使用 30% 的 GPU 記憶體
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.3)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

# 設定 Keras 使用的 TensorFlow Session
tf.keras.backend.set_session(sess)

# 使用 Keras 建立模型
# ...

自動增長 GPU 記憶體用量

直接設定 GPU 記憶體的用量可以確保程式不會吃掉過多的記憶體,但是如果程式遇到真的需要更大的記憶體時,就會因為記憶體不足而產生 ResourceExhaustedError,比較折衷的做法是採用自動增長 GPU 記憶體用量的方式,讓程式需要多少記憶體就拿多少,剩下的才留給別人:

import tensorflow as tf
W = tf.constant([1.0, 2.0, 3.0, 4.0], shape=[2, 2], name='W')
x = tf.constant([1.3, 2.4], shape=[2, 1], name='x')
y = tf.matmul(W, x)

# 自動增長 GPU 記憶體用量
gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
print(sess.run(y))

在 Keras 的部分也是一樣的做法:

import tensorflow as tf

# 自動增長 GPU 記憶體用量
gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

# 設定 Keras 使用的 Session
tf.keras.backend.set_session(sess)

# 使用 Keras 建立模型
# ...

參考資料:Kevin Chan’s blogcsdnStackOverflow