Building Tensorflow + CUDA + cuDNN

build_tf.sh

# Build Tensorflow + CUDA on ubuntu 16.04
#
# Modified from:
#  - http://expressionflow.com/2016/10/09/installing-tensorflow-on-an-aws-ec2-p2-gpu-instance/
#  - https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/docker/Dockerfile.devel-gpu
#  - https://github.com/NVIDIA/nvidia-docker/blob/master/ubuntu-16.04/cuda/8.0/devel/cudnn5/Dockerfile
#  - https://github.com/NVIDIA/nvidia-docker/blob/master/ubuntu-16.04/cuda/8.0/devel/Dockerfile
#  - https://github.com/NVIDIA/nvidia-docker/blob/master/ubuntu-16.04/cuda/8.0/runtime/Dockerfile

# Install base packages
apt-get update -y
apt-get upgrade -y
apt-get install -y \
  build-essential \
  git \
  python-pip \
  libfreetype6-dev \
  libxft-dev \
  libncurses-dev \
  libopenblas-dev \
  gfortran \
  python-matplotlib \
  libblas-dev \
  liblapack-dev \
  libatlas-base-dev \
  python-dev \
  python-pydot \
  libpng12-dev \
  libzmq3-dev \
  linux-headers-generic \
  linux-image-extra-virtual \
  unzip \
  python-numpy \
  swig \
  python-pandas \
  python-sklearn \
  unzip \
  wget \
  pkg-config \
  zip \
  g++ \
  zlib1g-dev \
  zip \
  libcurl3-dev \
  rsync \
  software-properties-common

# Install python
curl -fSsL -O https://bootstrap.pypa.io/get-pip.py && \
  python get-pip.py && \
  rm get-pip.py

pip --no-cache-dir install \
  ipykernel \
  jupyter \
  matplotlib \
  numpy \
  scipy \
  sklearn \
  && \
  python -m ipykernel.kernelspec

# Install CUDA + CUDNN
mkdir -p /root/build && cd /root/build

wget https://developer.nvidia.com/compute/cuda/8.0/prod/local_installers/cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64-deb && \
  dpkg -i cuda-repo-ubuntu1604-8-0-local_8.0.44-1_amd64-deb && \
  apt-get update && \
  apt-get install -y cuda

wget https://s3-us-west-2.amazonaws.com/vincentc/build/libcudnn6_6.0.21-1%2Bcuda8.0_amd64.deb && \
  dpkg -i libcudnn6_6.0.21-1+cuda8.0_amd64.deb

wget https://s3-us-west-2.amazonaws.com/vincentc/build/libcudnn6-dev_6.0.21-1%2Bcuda8.0_amd64.deb && \
  dpkg -i libcudnn6-dev_6.0.21-1+cuda8.0_amd64.deb

ln -s /usr/include/cudnn.h /usr/local/cuda/include/cudnn.h
ln -s /usr/lib/x86_64-linux-gnu/libcudnn.so /usr/local/cuda/lib64/libcudnn.so
ln -s /usr/lib/x86_64-linux-gnu/libcudnn.so.6.0.21 /usr/local/cuda/lib64/libcudnn.so.6

# Install Java8 (dependency of Bazel) + Bazel
add-apt-repository -y ppa:openjdk-r/ppa && \
  apt-get update && \
  apt-get install -y --no-install-recommends openjdk-8-jdk openjdk-8-jre-headless && \
  apt-get clean && \
  rm -rf /var/lib/apt/lists/*

echo "startup --batch" >>/root/.bazelrc
echo "build --spawn_strategy=standalone --genrule_strategy=standalone" >>/root/.bazelrc

export BAZELRC=/root/.bazelrc
export BAZEL_VERSION=0.8.1

cd /

mkdir /bazel && \
  cd /bazel && \
  curl -fSsL -O https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \
  chmod +x bazel-*.sh && \
  ./bazel-$BAZEL_VERSION-installer-linux-x86_64.sh && \
  cd /

git clone https://github.com/tensorflow/tensorflow.git && \
  cd tensorflow && \
  git checkout v1.4.0

# Build Tensorflow
export CI_BUILD_PYTHON=python
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/include:$LD_LIBRARY_PATH
export TF_NEED_CUDA=1
export TF_CUDA_COMPUTE_CAPABILITIES='3.0,3.5,5.2'
export CUDA_HOME=/usr/local/cuda
export CUDNN_INSTALL_PATH=/usr/local/cuda

tensorflow/tools/ci_build/builds/configured GPU \
  bazel build -c opt --config=cuda --config=opt tensorflow/tools/pip_package:build_pip_package && \
  bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/pip && \
  pip --no-cache-dir install --upgrade /tmp/pip/tensorflow-*.whl && \
  rm -rf /tmp/pip && \
  rm -rf /root/.cache

A few notes.

• tf version is v1.4.0
• Need to switch to bazel version 0.8.1
• need to update to libcudnn6_6.0.21-1+cuda8.0_amd64.deb and libcudnn6-dev_6.0.21-1+cuda8.0_amd64.deb
• symlink to cuDNN v6