TensorFlow 與 Keras 指定 NVIDIA GPU 顯示卡與記憶體用量教學

本篇介紹如何指定 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 blog、csdn、StackOverflow