環境介紹
作業系統Ubuntu 16.04.5,其中docker 18.09.0已安裝,nvidia顯示卡驅動已安裝。
步驟
(1)安裝nvidia-docker
nvidia-docker是docker引擎的一個應用插件,專門面向nvidia的GPU。官方github給出的nvidia-docker在Ubuntu 16.04中的安裝方法如下:
#If you have nvidia-docker 1.0 installed: we need to remove it and all existing GPU containers
docker volume ls -q -f driver=nvidia-docker|xargs -r -I{} -n1 docker ps -q -a -f volume={}|xargs -r docker rm -f
sudo apt-get purge -y nvidia-docker
# Add the package repositories
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey |
sudo apt-key add -
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
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
# Install nvidia-docker2 and reload the Docker daemon configuration
sudo apt-get install -y nvidia-docker2
sudo pkill -SIGHUP docker
測試安裝是否成功:
# Test nvidia-smi with the latest official CUDA image
docker run --runtime=nvidia --rm nvidia/cuda:9.0-base nvidia-smi
如果nvidia-docker安裝成功,系統會顯示GPU資訊:

(2)啟動pytorch容器
拉取合适的pytorch鏡像:
docker pull floydhub/pytorch
通過nvidia-docker啟動容器,容器名稱為torch,容器内目錄/workspace挂載于伺服器目錄~/leon/pytorch:
nvidia-docker run -it -d --name="torch" -v ~/leon/pytorch:/workspace pytorch/pytorch:latest
以互動模式進入容器:
docker exec -it torch /bin/bash
進入python控制台,可以通過pytorch在docker内使用nvidia顯示卡: