docker容器中执行GPU环境中的tensorflow和pytorch任务

1. 背景

(1) 业务方提供了一台有GPU的服务器，且已经安装了显卡等组件，cuda版本10.2，具体信息如下

 

(2) 在裸机上部署anaconda、pytorch、tensorflow较为麻烦，因此打算在docker中执行pytorch和tensorflow任务

2. 部署及使用

2.1 先决条件

1) 必须要NVIDIA显卡

2) 安装NVIDIA显卡驱动，可参CentOS 7 安装 NVIDIA 显卡驱动和 CUDA Toolkit，https://github.com/NVIDIA/nvidia-docker，nvidia驱动下载

3) 安装docker，版本最好>=19.*(本人安装的是18.*)

2.2 安装GPU版本的tensorflow

1) docker运行GPU版的tensorflow，只需要安装NVIDIA驱动即可，而不需要安装cuda相关的工具包，docker镜像中已经包含

2) 参考nvidia-docker，tensorflow docker安装 

2.3 使用GPU版本的tensorflow

1)  docker中测试当前安装的nvidia驱动

docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi  

2) 测试tensorflow的GPU设备信息

docker run --runtime=nvidia -it --rm tensorflow/tensorflow:latest-gpu   python -c "import tensorflow as tf; print(tf.test.gpu_device_name())" 

3) docker容器中执行命令

docker run -it tensorflow/tensorflow:latest-gpu bash 

4) 容器中执行tensorflow 脚本

docker run -it --rm -v $PWD:/tmp -w /tmp tensorflow/tensorflow:latest-gpu python ./script.py 

2.4 安装GPU版本的pytorch

拉取镜像：nvcr.io/nvidia/pytorch:latest

2.5 使用GPU版本的pytorch

1) docker容器中执行命令

docker run --runtime=nvidia -v ${PWD}:/tmp  -it --rm nvcr.io/nvidia/pytorch:20.03-py3 bash