首页 > 解决方案 > 如何在 ubuntu 16.04 上安装 Tensorflow GPU

问题描述

我正在使用带有 Tesla K80 GPU 和 Ubuntu 16.04 LTS 的 AWS P2 gpu 实例。为 CPU 安装 Tensorflow 很简单:

pip install tensorflow

但是为 GPU 安装 Tensorflow 不起作用:

pip install tensorflow-gpu

标签: ubuntutensorflownvidia

解决方案


安装英伟达驱动:

sudo apt-get remove nvidia* && sudo apt autoremove  
sudo reboot  # you will have to wait till reboot completes
sudo apt-get update
sudo apt-get install openjdk-8-jdk
curl -O http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-repo-ubuntu1604_9.0.176-1_amd64.deb
sudo apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/7fa2af80.pub
sudo dpkg -i ./cuda-repo-ubuntu1604_9.0.176-1_amd64.deb
sudo apt-get update
sudo apt-get install cuda-9-0  

以下代码应向您显示有关您的 gpu 的详细信息

nvidia-smi

如果您看到有关 GPU 的详细信息,请继续执行以下命令

wget https://s3.amazonaws.com/open-source-william-falcon/cudnn-9.0-linux-x64-v7.1.tgz  
sudo tar -xzvf cudnn-9.0-linux-x64-v7.1.tgz  
sudo cp cuda/include/cudnn.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn.h /usr/local/cuda/lib64/libcudnn*
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
export CUDA_HOME=/usr/local/cuda
source ~/.bashrc

确保卸载之前的 TensorFlow。然后激活您的虚拟环境并在其中安装 tensorflow_gpu。如果您没有虚拟环境(您应该),以下代码将为您提供帮助:

sudo pip install virtualenv
mkdir ~/virtualenv #creating a folder to store all your virtual env
cd ~/virtualenv/
virtualenv -p python3 venv1 #creates venv1 with new python installed on your ubuntu
source venv1/bin/activate #activates your venv
pip uninstall tensorflow
pip uninstall tensorflow-gpu
pip install jupyter # coz i like jupyter
pip install tensorflow-gpu

现在让我们测试一下 Tensorflow-gpu 是否有效

python
import tensorflow as tf   
hello = tf.constant('Hello, TensorFlow!') #if no error received, its working
from tensorflow.python.client import device_lib
device_lib.list_local_devices() #this will show all CPU's and GPU's on your system

耶!!


推荐阅读