在docker容器中使用tensorflow-gpu

目前的项目需要在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一下即可测试是否成功。

最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
【社区内容提示】社区部分内容疑似由AI辅助生成,浏览时请结合常识与多方信息审慎甄别。
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

友情链接更多精彩内容