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

1. 背景

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

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

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

2. 部署及使用

2.1 先决条件

1) 必须要NVIDIA显卡

2) 安装NVIDIA显卡驱动,可参CentOS 7 安装 NVIDIA 显卡驱动和 CUDA Toolkithttps://github.com/NVIDIA/nvidia-dockernvidia驱动下载

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

2.2 安装GPU版本的tensorflow

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

2) 参考nvidia-dockertensorflow 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
docker容器中执行GPU环境中的tensorflow和pytorch任务

 

上一篇:ubuntu16.04安装NVIDIA显卡驱动相关问题


下一篇:win10安装Cuda环境