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