Ubuntu安装GPU版TensorFlow

文章目录
  1. 环境准备
  2. 安装CUDA和cuDNN
  3. 安装TensorFlow

2017.12.5 CUDA 与 cuDNN的最新版本均已经有所变化,不过安装方式相同。

当前环境如下:

  • Ubuntu 16.04.1 LTS
  • GeForce GTX 960M
  • CUDA 8.0
  • cuDNN V5.1 Library

环境准备

首先检查自己的GPU是否支持CUDA,

1
2
#查看显卡类型
$ lspci | grep -i nvidia

在输出中找到自己的的显卡类型,然后到http://developer.nvidia.com/cuda-gpu 找到自己的型号:

我本机的环境如下:

  • Ubuntu 16.04.1 LTS
  • GeForce GTX 960M

下面开始安装。

安装CUDA和cuDNN

  1. 下载CUDA: https://developer.nvidia.com/cuda-downloads,下载对应版本的cuda安装包,下载对应版本的cuda安装包,推荐下载最新版本的8.0版本,deb版的本地安装比较方便。

  2. 下载完成之后,直接执行如下的命令进行安装即可:

    1
    2
    3
    sudo dpkg -i cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64.deb
    sudo apt update
    sudo apt install cuda

    注: 下载完成之后,deb的名称可能有一个oauth的后缀,直接重命令删掉就好了。

  3. 下载cuDDN:https://developer.nvidia.com/rdp/cudnn-download , 这里需要注册个帐号,然后才可以下载。

    这里有个版本需要注意,最好下载 cuDNNv5.1 Library版本的。 虽然V5.0也支持 CUDA 8.0,但是在后续跑TensorFlow的时候,会报错,后来才重新下载了V5.1的,才能正常运行。

  4. 拷贝库文件 将下载下的cuDNN压缩包执行如下命令:

    1
    2
    3
    4
    tar -zxvf cudnn-8.0-linux-x64-v5.1.tgz
    cd cuda
    sudo cp lib64/* /usr/local/cuda/lib64/ -d
    sudo cp include/cudnn.h /usr/local/cuda/include/

  5. 修改环境变量 修改.bashrc或者/etc/profile等,加入如下的环境变量:

    1
    2
    3
    4
    ## CUDA
    export CUDA_HOME=/usr/local/cuda
    export PATH=${CUDA_HOME}/bin/:$PATH
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_HOME/lib64:$CUDA_HOME/extras/CUPTI/lib64

    之后执行source .bashrc立即生效

  6. 测试是否成功 进入/usr/local/cuda/samples,编译运行某个sample:

    1
    sudo make all -j4

    全部编译完成后,进入 samples/bin/x86_64/linux/release,运行deviceQuery

    1
    ./deviceQuery

    显示有可用的GPU,即可成功:

安装TensorFlow

直接使用pip安装即可:当前版本为1.4:

1
sudo pip install tensorflow-gpu

安装完成之后,尝试import tensorflow:显示如下即可成功(最新版tf显示信息有所差异)

1
2
3
4
5
6
7
In [1]: import tensorflow
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcublas.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcudnn.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcufft.so locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcuda.so.1 locally
I tensorflow/stream_executor/dso_loader.cc:128] successfully opened CUDA library libcurand.so locally
In [2]:

在代码使用的时候,只需要指定GPU即可(当然tf也会默认使用GPU):

1
with tf.device('/gpu:0'):

即可使用GPU运行TensorFlow程序。速度比CPU快很多,跑一个CNN文本分类的程序,CPU版本,用8核i7处理器几乎保持600%的利用率,训练半小时还未结束,换为GPU版本,10分钟结束! 另外提供一个使用watch实时查看nvidia显卡状态的命令,每隔1s显示状态:

1
watch -n 1 nvidia-smi