在Ubuntu使用nvidia-docker从零搭建pytorch容器环境

0. 搭建环境

Ubuntu 22.04 LTS desktop版

下面所有的命令都是在root用户下运行的，所以请先获得管理员权限。

1. 安装docker

a) 设置仓库

apt update # 更新源
apt install ca-certificates curl gnupg lsb-release # 安装依赖
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg # 将官方Docker存储库的GPG密钥添加到系统
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null # 设置稳定的仓库。

b) 安装docker-ce和配套

apt update # 上一步设置仓库后，要更新源
apt install docker-ce docker-ce-cli containerd.io docker-compose-plugin # 安装docker-ce和配套

2. 安装nvidia驱动

ubuntu-drivers devices # 查看显卡型号和推荐的驱动
ubuntu-drivers autoinstall # 自动安装推荐的驱动

如果对驱动版本有特殊的需求，apt install 需要的版本即可。

3. 安装nvidia-docker2

设置仓库和GPG密钥

distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
      && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
      && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
            sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

更新源并安装nvidia-docker

apt update # 更新源
apt install -y nvidia-docker2 # 安装nvidia-docker
systemctl restart docker # 重启docker

上DockerHub查需要的pytorch/pytorch镜像版本，如果没有特定版本需求，拉pytorch/pytorch:latest即可

docker pull pytorch/pytorch[:tag]

直接使用这个镜像测试一下，容器内能正常打印nvidia-smi就OK了。

docker run -rm --gpus all pytorch/pytorch nvidia-smi

4. docker换源

安装nvidia-docker2会覆盖docker的daemon文件，所以之前换过源的需要重新写入一次daemon文件。在/etc/docker/daemon.json后追加下面内容

{
  "registry-mirrors": [
    "https://registry.docker-cn.com",
    "https://docker.mirrors.ustc.edu.cn",
    "https://hub-mirror.c.163.com"
  ]
}

5. 修改docker的镜像存储地址到机械硬盘上

docker容器和镜像的存储地址默认是/var/lib/docker,这大概率是在系统盘上。如果你的系统盘不是特别大，最好把默认存储地址设置到机械硬盘上。

docker info | grep "Docker Root Dir"  # 查看docker对象存储位置
# 基本默认都是/var/lib/docker

备份已有镜像（如果你有的话）

docker save -o filename.tar image_name

以软连接的方式更改 Docker 镜像存储路径

# 停止 docker
service docker stop 

# 移动到新路径
mv /var/lib/docker /某个机械硬盘上的地址
# 创建软连接
ln -s /某个机械硬盘上的地址/docker /var/lib/docker 

# 启动 docker
service docker start 

或者可以以修改配置文件的方式更新 docker 的默认镜像存储路径

# 停止 docker
service docker stop

# 编辑文件 docker-overlay.conf
cd /etc/systemd/system/docker.service.d # 如果没有docker.service.d 则创建该路径
vim docker-overlay.conf  # 如果没有则创建该文件

在文件中添加如下内容：

[Service]
ExecStart=/usr/bin/dockerd --graph=/某个机械硬盘上的地址/docker --storage-driver=overlay

重新启动 docker

systemctl daemon-reload # 重新载入daemon
service docker start 
docker info | grep "Docker Root Dir" # 检查保存地址
reboot # 重启
docker info | grep "Docker Root Dir" # 再检查一次保存地址

6. 打开一个pytorch容器

# 拉取pytorch镜像
docker pull pytorch/pytorch[:tag] # 或者换成你需要的版本的pytorch镜像
# 开一个样板镜像，使用nvidia-docker最重要的参数是--gpus
docker run --gpus all -it --name test pytorch/pytorch /bin/bash

然后你就进入名字为test的容器内了。

容器内置配好pytorch的python3在/opt/conda/bin/python。

# 运行python
/opt/conda/bin/python
>> import torch
>> print(torch.__version__) # 查看pytorch版本
>> print(torch.cuda.is_available()) # 查看cuda是否可用

如果正常打印，应该就OK了，愉快使用吧 😃