1. 程式人生 > >【tensorflow 學習】 gpu使用

【tensorflow 學習】 gpu使用

由於tensorflow預設搶佔伺服器所有GPU視訊記憶體,只允許一個小記憶體的程式也會佔用所有GPU資源。下面提出使用GPU執行tensorflow的幾點建議:

1.在執行之前先檢視GPU的使用情況:

$ nvidia-smi # 檢視GPU此時的使用情況

或者

$ nvidia-smi -l # 實時返回GPU使用情況

2.目前實驗室伺服器有0,1,2,3四個GPU,找到空閒的GPU號,可以使用環境變數CUDA_VISIBLE_DEVICES

環境變數的定義格式: 
CUDA_VISIBLE_DEVICES=1 Only device 1 will be seen 
CUDA_VISIBLE_DEVICES=0,1 Devices 0 and 1 will be visible 
CUDA_VISIBLE_DEVICES=”0,1” Same as above, quotation marks are optional 
CUDA_VISIBLE_DEVICES=0,2,3 Devices 0, 2, 3 will be visible; device 1 is masked

輸入以下命令執行程式:

$ export CUDA_VISIBLE_DEVICES=0 # 假設此時 GPU 0 空閒

為了防止新開終端忘記export,比較保險的做法是每次執行tensorflow之前定義使用的GPU:

$ CUDA_VISIBLE_DEVICES=0 python mnist.py # 假設此時 GPU 0 空閒, mnist.py為你想執行的程式。

3.這樣tensorflow此時只會在指定的GPU上執行,但是仍然會佔用整個GPU的視訊記憶體,不過不和其他人公用GPU時也不會有影響,下面介紹兩種限定GPU佔用的方法:

(1)在tensorflow中定義session時作如下設定,該設定會啟用最少的GPU視訊記憶體

來執行程式。

config = tf.ConfigProto() 
config.gpu_options.allow_growth = True 
session = tf.Session(config=config) 

(2)在tensorflow中定義session時作如下設定,該設定會強制程式只佔用指定比例的GPU視訊記憶體

config = tf.ConfigProto() 
config.gpu_options.per_process_gpu_memory_fraction = 0.4 # 佔用GPU40%的視訊記憶體 
session = tf.Session(config=config)

注: 
- 在tensorflow程式碼中with tf.device(‘/gpu:0’):只會指定在GPU 0上計算,但仍然會預設佔用所有GPU資源。

4.檢視使用的裝置 
tf.ConfigProto(log_device_placement=True))

# Creates a graph.
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print(sess.run(c))

執行後能看到如下結果:

Device mapping:
/job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: Tesla K40c, pci bus
id: 0000:05:00.0
b: /job:localhost/replica:0/task:0/device:GPU:0
a: /job:localhost/replica:0/task:0/device:GPU:0
MatMul: /job:localhost/replica:0/task:0/device:GPU:0
[[ 22.  28.]
 [ 49.  64.]]