TensorFlow 指定 GPU 訓練模型

Linux 查看當前服務器 GPU 的佔用情況可以使用 nvidia-smi 命令，如下所示：

nvidia-smi

關於 nvidia-smi 命令輸出的詳細解釋，可參考筆者的另外一篇博客：GPU狀態監測 nvidia-smi 命令詳解。在此不再贅述，本文主要分享一下在用 TensorFlow 訓練模型時如何指定 GPU。

在用 TensorFlow 訓練深度學習模型的時候，若無手動指定，默認是選用第0塊來訓練，而且其他幾塊 GPU 也會被顯示佔用。有時候，我們更希望可以自己指定一塊或者多塊 GPU 來訓練模型，接下來介紹一種常用的也是比較有效的方法：

舉個例子，如果要指定只用第0塊 GPU 訓練，可以在 python 代碼中如下指定：

import os
os.environ['CUDA_VISIBLE_DEVICES']='0'

運行 python 程序後，可發現就只有第0塊 GPU 顯示被佔用，如上圖所示。

如果要指定多塊 GPU，可以如下：

os.environ['CUDA_VISIBLE_DEVICES']='1,2'

當然，上述這種方法不太方便，每次運行若換用 GPU 訓練都需要改代碼，可以在運行 python 程序的時候進行指定：

CUDA_VISIBLE_DEVICES=0 python textCnn.py