TensorFlow GPU 的使用

1、TensorFlow 設備分配

一、設備分配規則

If a TensorFlow operation has both CPU and GPU implementations, the GPU devices will be given priority when the operation is assigned to a device.session

二、手動指定設備分配

  • 若是你不想讓系統自動爲 operation 分配設備, 而是本身手動指定, 能夠用 with tf.device 建立一個設備環境, 這個環境下的 operation 都統一運行在指定的設備上.
  • 代碼示例以下:
 1 # op 在 cpu 上運算
 2 with tf.device('/cpu:0'):
 3       a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a') 4 b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b') 5 6 # op 在 gpu 上運算 7 with tf.device('/device:GPU:2'): 8 a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a') 9 b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b') 10 11 # op 在 gpus 上運算 12 for d in ['/device:GPU:2', '/device:GPU:3']: 13  with tf.device(d): 14 a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3]) 15 b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])

2、TensorFlow GPU 配置

一、指定能夠被看見的GPU設備spa

1 import os
2 
3 # 默認狀況,TF 會佔用全部 GPU 的全部內存, 咱們能夠指定
4 # 只有 GPU0 和 GPU1 這兩塊卡被看到,從而達到限制其使用全部GPU的目的
5 os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'  
6 
7 # 打印 TF 可用的 GPU
8 print os.environ['CUDA_VISIBLE_DEVICES'] 9 >>> 0, 1

 二、限定使用顯存的比例code

 1 # 在開啓對話session前,先建立一個 tf.ConfigProto() 實例對象
 2 # 經過 allow_soft_placement 參數自動將沒法放在 GPU 上的操做放回 CPU
 3 gpuConfig = tf.ConfigProto(allow_soft_placement=True)
 4 
 5 # 限制一個進程使用 60% 的顯存
 6 gpuConfig.gpu_options.per_process_gpu_memory_fraction = 0.6
 7 
 8 # 把你的配置部署到session
 9 with tf.Session(config=gpuConfig) as sess: 10 pass 11 12 這樣,若是你指定的卡的顯存是8000M的話,你這個進程只能用4800M。

三、須要多少拿多少

 1 # 在開啓對話session前,先建立一個 tf.ConfigProto() 實例對象
 2 # 經過 allow_soft_placement 參數自動將沒法放在 GPU 上的操做放回 CPU
 3 gpuConfig = tf.ConfigProto(allow_soft_placement=True)
 4 
 5 # 運行時須要多少再給多少
 6 gpuConfig.gpu_options.allow_growth = True 7 8 # 把你的配置部署到session 9 with tf.Session(config=gpuConfig) as sess: 10 pass

  四、GPU 使用總結

1 import os
2 os.environ['CUDA_VISIBLE_DEVICES'] = '0, 1'  
3 
4 gpuConfig = tf.ConfigProto(allow_soft_placement=True) 5 gpuConfig.gpu_options.allow_growth = True 6 7 with tf.Session(config=gpuConfig) as sess: 8 pass
相關文章
相關標籤/搜索