Tensorflow-gpu Docker 环境搭建 (一)

发表于 分类于 阅读次数： Disqus： 本文字数： 3.4k 阅读时长 ≈ 3 分钟

  • 构建 Tensorflow 支持 GPU Docker 镜像

  • 资料来源:

    https://www.tensorflow.org/install/docker?hl=zh-cn
    https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/user-guide.html
    https://github.com/andreivmaksimov/python_data_science

  • 更新

    1
    20.10.12 初始化

导语

Win 下的 TF 环境运行不错,没出啥问题…但是主力是 Ubuntu,总归需要在 Linux 下安装这一套环境.

TF+GPU 支持要安装一堆包,按照习惯是上 Docker.不过这个坑有点大,分成两期了.

这一篇仅仅是将 Docker 弄到能运行的程度.

特别参考了 GitHub - andreivmaksimov/python_data_science

需求

  • 可以运行常见深度学习工具包,特别是 Sklearn 和 Keras.
  • 支持 GPU 加速
  • 对 python 环境的改动不需要重新创建容器.

启用 Dockr GPU 支持

在 Docker >= 19.03 版本以后,Docker 内 GPU 支持非常简单

  • 宿主机安装 Nvdia 专有驱动和 nvidia-container-toolkit.
  • 启动镜像时带有 --gpus all 参数.
  • 注意: 一些教程中安装 nvidia-docker2,目前官方已不再更新(但是还能用),及时切换到新版吧.

安装 nvidia-container-toolkit ,Nvida 官方最坑的就是这,在 Installation Guide 居然还是安装 nvidia-docker2..

执行下面的命令

1
2
3
4
5
6
7
8
9
10
# 添加环境变量
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
# 安装 nvidia-container-toolkit
sudo apt-get install nvidia-container-toolkit -y

之后重启 docker 验证

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
sudo systemctl restart docker
# 验证
docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi

# 正常输出 GPU 信息
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.51.06    Driver Version: 450.51.06    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   34C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

顺带说一下 nvidia/cuda 提供的 3 种镜像(虽然用不到了)

  • base: 最基本的运行环境
  • runtime: 多了一些运行时的库
  • devel: 最完整.

DockerFile

启用 GPU 支持,TF - CUDA - 驱动版本必须是对应的.

  • CUDA - 驱动对照表
  • TF - CUDA 对照表

目前 Ubuntu 稳定版驱动是 450,也就是最高支持到支持 CUDA 11.0.同时这个版本也支持 TF 2.2,这里就维持这个版本对照了.

尝试了 N 多 nvidia-docker 镜像,一直遇到 pip 安装了 tf 还是提示有关的库不全.

TF 官方已经提供了基础镜像 Docker Hub,这里选择 tensorflow/tensorflow:2.2.1-gpu-py3-jupyter 作为基础镜像.

其他的库都通过 pip 进行安装.对 jupyter 单独使用还不算熟悉,未来或许会精简一些.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
FROM tensorflow/tensorflow:2.2.1-gpu-py3-jupyter

ENV DEBIAN_FRONTEND noninteractive

RUN pip3 install numpy pandas sklearn matplotlib seaborn pyyaml h5py && \
    pip3 install keras --no-deps && \
    pip3 install opencv-python && \
    pip3 install imutils

RUN ["mkdir", "notebooks"]
COPY conf/.jupyter /root/.jupyter
COPY run_jupyter.sh /

# Jupyter and Tensorboard ports
EXPOSE 8888 6006

# Store notebooks in this mounted directory
VOLUME /notebooks

CMD ["/run_jupyter.sh"]

其他文件见 GitHub - andreivmaksimov/python_data_science,未改动.

之后正常 build ,启动镜像千万不要忘了附加 --gpus all 参数.

测试

测试的脚本还是之前用过的 test.py

1
2
3
4
5
6
7
8
9
10
#!/usr/bin/python3

import tensorflow as tf

a = tf.test.is_built_with_cuda()  # 判断CUDA是否可以用
b = tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None) # GPU 支持
c = tf.test.gpu_device_name() # GPU 设备
print(a)
print(b)
print(c)

在容器内运行,得到俩 true 和 GPU 信息

结语

被前面的 nvidia-cuda 坑了好久,不知道为什么一直提示缺少库.

整个的镜像还是蛮大的 4.36G..原始的 TF 镜像已经 3.96G 了..

下一步打算把虚拟的 python 环境也打包进去,省得每次都需要重新编译镜像.

相关文章