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