docker挂载NVIDIA显卡运行pytorch的方法

2023年5月25日下午1:09 • 人工智能概览

下面我将详细讲解"docker挂载NVIDIA显卡运行pytorch的方法"。

1. 安装NVIDIA驱动和docker

首先，我们需要在宿主机上安装NVIDIA的显卡驱动，以及在宿主机上安装docker。关于这两个软件的安装过程这里不再赘述，如果你还没有安装，请自行搜索相关教程。

2. 下载nvidia/cuda镜像

使用以下命令下载nvidia/cuda镜像：

docker pull nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04

其中，10.0是CUDA版本号，cudnn7是cudnn的版本号，ubuntu18.04是基础镜像的操作系统版本号。你可以根据自己的需求选择不同的版本号。

3. 启动nvidia-docker2服务

接下来，我们需要启动nvidia-docker2服务：

sudo systemctl start nvidia-docker2

4. 启动带有GPU支持的docker容器

我们可以使用以下命令启动docker容器：

docker run --gpus all -it --rm nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04 /bin/bash

其中，--gpus all表示将所有NVIDIA显卡都挂载到docker容器中，-it表示交互模式运行docker容器，--rm表示当docker容器退出后会自动删除容器，nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04是下载的镜像名称，/bin/bash表示在容器中启动bash终端。

5. 在docker容器中安装pytorch和torchvision

在docker容器中安装pytorch和torchvision，我们可以使用以下命令：

pip install torch torchvision

示例1：在docker容器中运行pytorch代码

下面是一个简单的示例，演示如何在docker容器中运行pytorch代码。首先，我们需要先在宿主机上创建一个名为test.py的文件，文件内容如下：

import torch

x = torch.randn(1)
print(x)

接下来，在宿主机上打开一个终端窗口，输入以下命令：

docker run --gpus all -v $(pwd):/work -it --rm nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04 /bin/bash

其中，-v $(pwd):/work表示将宿主机的当前目录挂载到docker容器的/work目录中。

然后，在docker容器中，使用以下命令运行test.py：

python /work/test.py

执行完毕后，输出结果如下：

tensor([0.5476])

示例2：在docker容器中训练MNIST数据集的模型

下面是一个稍微复杂的示例，演示如何在docker容器中使用pytorch训练MNIST数据集的模型。首先，我们需要先在宿主机上创建一个名为train_mnist.py的文件，文件内容如下：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 超参数
batch_size = 64
learning_rate = 0.01
num_epochs = 10

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载MNIST数据集
train_dataset = datasets.MNIST(root='data/', train=True, transform=transforms.ToTensor(), download=True)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

test_dataset = datasets.MNIST(root='data/', train=False, transform=transforms.ToTensor(), download=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=True)

# 定义模型
model = nn.Sequential(
    nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2, stride=2),
    nn.Flatten(),
    nn.Linear(32 * 7 * 7, 10),
    nn.Softmax(dim=1)
)

model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# 训练模型
for epoch in range(num_epochs):
    for batch_idx, (data, targets) in enumerate(train_loader):
        data = data.to(device=device)
        targets = targets.to(device=device)

        # 前向传播
        scores = model(data)
        loss = criterion(scores, targets)

        # 反向传播并优化参数
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 测试模型
def check_accuracy(loader, model, device):
    num_correct = 0
    num_samples = 0
    model.eval()

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device=device)
            y = y.to(device=device)

            scores = model(x)
            _, predictions = scores.max(1)
            num_correct += (predictions == y).sum()
            num_samples += predictions.size(0)

        print(f'Got {num_correct}/{num_samples} with accuracy {float(num_correct)/float(num_samples)*100:.2f}')

    model.train()

check_accuracy(train_loader, model, device)
check_accuracy(test_loader, model, device)

然后，在宿主机上执行以下命令：

docker run --gpus all -v $(pwd):/work -it --rm nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04 /bin/bash

然后，在docker容器中，使用以下命令运行train_mnist.py：

python /work/train_mnist.py