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

0. 搭建环境

Ubuntu 22.04 LTS desktop版

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

1. 安装docker

a) 设置仓库

1
2
3
4
5
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和配套

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

2. 安装nvidia驱动

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

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

3. 安装nvidia-docker2

设置仓库和GPG密钥

1
2
3
4
5
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

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

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

1
docker pull pytorch/pytorch[:tag]

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

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

4. docker换源

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

1
2
3
4
5
6
7
{
"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,这大概率是在系统盘上。如果你的系统盘不是特别大,最好把默认存储地址设置到机械硬盘上。

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

备份已有镜像(如果你有的话)

1
docker save -o filename.tar image_name

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

1
2
3
4
5
6
7
8
9
10
# 停止 docker
service docker stop

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

# 启动 docker
service docker start

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

1
2
3
4
5
6
# 停止 docker
service docker stop

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

在文件中添加如下内容:

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

重新启动 docker

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

6. 打开一个pytorch容器

1
2
3
4
# 拉取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

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

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