目前的项目需要在docker容器中用到tensorflow-gpu,项目需要用的是编译版的tensorflow,在这过程中踩了不少坑,写下本文总结一下全过程。
0.环境说明
主机OS是Ubuntu18.04,已经安装了docker和nvidia显卡驱动,如何安装不再赘述了。
1.安装Nvidia-Container组件
除了显卡驱动,我们还需要一个Nvidia-Container组件来使docker容器内也可以使用主机的显卡。
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
$ curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
$ curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
$ sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
$ sudo systemctl restart docker
安装完成后,即可输入以下代码来测试组件是否安装成功。
$sudo docker run --gpus all nvidia/cuda:9.0-base nvidia-smi

如图所示的情况即为成功。
2.启动docker容器
为了节省下载tensorflow包的时间,这里选用了tensorflow的开发者镜像,该镜像内自带tensorflow和cuda,可以直接编译安装完整版tensorflow(但还是很慢)。启动前最好设置一下docker默认大小,等生成容器了再改就不太方便了,此部分将在另一篇中描述如何操作。
docker run --gpus all -p 8080:8899 tensorflow/tensorflow:devel-gpu
为了能够在另一台电脑能够使用tensorflow,所以在启动时指定了8080端口来对接容器内8899端口。
进入到docker容器内,首先要做的当然就是更新
apt update -y
apt upgrade -y
upgrade那一步因为cudnn的原因可能要更新很久,大概两小时,有能力的同学可以自己下cudnn文件放到容器中省去更新(不过也快不了多少)。
更新后要安装bazel,利用bazel来编译tensorflow。
sudo aptinstall
curl gnupgcurl-fsSL https://bazel.build/bazel-release.pub.gpg | gpg--dearmor>bazel.gpg
sudo mv bazel.gpg /etc/apt/trusted.gpg.d/
echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | sudo tee /etc/apt/sources.list.d/bazel.list
apt install bazel
3.编译tensorflow
首先需要进入到tensorflow的路径中,对编译过程进行设置。
./configure
过程中大部分需要路径的都可以enter选择默认路径,其余的y/n的都选n。
设置完成后,运行下面的代码进行编译,这个过程大约在40分钟左右。
bazel build --config=cuda //tensorflow/tools/pip_package:build_pip_package
编译完成后到根目录的tmp文件夹中创建一个名为tensorflow_pkg的文件夹,再回到tensorflow的文件夹中,将安装包复制到tmp中。
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg
复制完成后,回到tmp文件夹中运行下方代码:
pip install /tmp/tensorflow_pkg/tensorflow-version-tags.whl
#version-tags可以用tab自动生成
最后进入python,import一下即可测试是否成功。