TensorFlow 在指定的 GPU 上運行及資源使用設置方法

一、TensorFlow 指定計算使用的 GPU 卡

• 終端運行命令時指定（推薦使用該方法）

~$ CUDA_VISIBLE_DEVICES=0 python file_name.py	# 使用第 0 號 GPU
~$ CUDA_VISIBLE_DEVICES=1 python file_name.py	# 使用第 1 號 GPU
~$ CUDA_VISIBLE_DEVICES=0,1,2 python file_name.py	# 使用第 1,2,3 號 GPU
~$ CUDA_VISIBLE_DEVICES="0,1,2" python file_name.py	# 使用第 1,2,3 號 GPU
~$ CUDA_VISIBLE_DEVICES="" python file_name.py # 不使用 GPU

• Python 代碼中指定 GPU

>>> import os
>>> os.environ["CUDA_VISIBLE_DEVICES"] = "0"	# 使用第 0 號 GPU
>>> os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"	# 使用第 1,2,3 號 GPU

二、TensorFlow 設置使用顯存的大小

• 方法一：設置 config 對象的屬性：gpu_options.per_process_gpu_memory_fraction

>>> config = tf.ConfigProto()
>>> config.gpu_options.per_process_gpu_mem>>> ory_fraction = 0.6	# 使用 60% 的顯存
>>> session = tf.Session(config=config)

• 方法二：創建 GPUOptions 時設置參數：per_process_gpu_memory_fraction

>>> gpu_options=tf.GPUOptions(per_process_gpu_memory_fraction=0.6)	# 使用 60% 的顯存
>>> config=tf.ConfigProto(gpu_options=gpu_options)
>>> session = tf.Session(config=config)

三、按需動態使用 GPU 資源

# 開始階段分配少量 GPU 資源，然後按需增加分配的 GPU 資源
>>> config = tf.ConfigProto()
>>> config.gpu_options.allow_growth = True
>>> session = tf.Session(config=config)

四、查看 GPU 使用情況

~$ nvidia-smi
# nvidia-smi 各種參數的詳細功能及用法可以通過： nvidia-smi -h 來查看，結合具體需要使用