CUDA 11.5 对应的 PyTorch 指南

CUDA 是 NVIDIA 公司推出的一个并行计算框架，它可为 GPU 提供良好的并行计算环境和编程接口，而 PyTorch 是一个基于 Python 的机器学习框架，这两个框架的结合可使我们更高效地进行深度学习开发。本篇文章将从多个方面对 CUDA 11.5 对应的 PyTorch 做详细的阐述。

一、 安装和配置

安装 CUDA 11.5 可以去官网下载安装。对于 Python 环境的设置，可以通过 Anaconda 来管理。首先，新建并激活一个名为 pytorch 的虚拟环境，然后通过下面代码安装 PyTorch。

conda create -n pytorch python=3.8
conda activate pytorch
conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch

注：cudatoolkit 的版本需要和本机安装的 CUDA 版本对应，这里是 11.1 是因为 PyTorch 官方没有在 11.5 上发布的相关版本，而在 11.1 上是可用的。

接着，在运行 PyTorch 代码时，需要注意以下几点：

1. 确认安装的 PyTorch 版本是否支持当前 CUDA 的版本；

2. 确认 GPU 是否支持 CUDA；

3. 通过下面的代码检查 PyTorch 是否安装成功。

import torch

print(torch.cuda.is_available())  # 检查 GPU 是否支持 CUDA
print(torch.__version__)          # 检查 PyTorch 版本号

二、 数据预处理

在进行深度学习训练时，数据预处理是一个非常重要的环节。PyTorch 提供了多种方式来进行数据预处理。对于图像数据，我们可以使用 torchvision 来预处理数据，并将其转换为 PyTorch 中的 Tensor 数据类型。

下面是一个示例代码，它使用了 torchvision.transforms 中的 Compose、Resize、ToTensor 函数，将图像数据转换为 PyTorch 中的 Tensor 数据类型。

from torchvision.transforms import Compose, Resize, ToTensor
from PIL import Image

dataset_dir = 'your dataset directory'
dataset = []

# 加载图像并进行预处理
transform = Compose([Resize((224, 224)), ToTensor()])

for img_path in os.listdir(dataset_dir):
    img = Image.open(os.path.join(dataset_dir, img_path)).convert('RGB')
    x = transform(img)
    dataset.append(x)

上述代码中，Componse 函数用于将多个预处理操作组合在一起，Resize 函数用于调整图像大小，ToTensor 函数用于将图像数据转换为 Tensor 数据类型。在循环中，我们使用 PIL 库来加载图像数据，将其转换为 PyTorch 中的 Tensor 数据类型，并将其添加到 dataset 列表中。

三、 神经网络搭建

在 PyTorch 中，我们可以使用 nn.Module 类来构建神经网络模型。该类提供了一些有用的函数，例如 forward 函数，用于定义输入和输出之间的计算。

下面是一个示例代码，它演示如何使用 nn.Module 类构建一个简单的神经网络，并使用该网络对图像进行分类。

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
        self.fc1 = nn.Linear(32 * 56 * 56, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 32 * 56 * 56)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

net = Net()

# 在 GPU 上运行模型
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
net.to(device)

上述代码中，我们首先定义了一个 Net 类，并在 __init__ 构造函数中定义了神经网络的结构，包括卷积层、池化层和全连接层。在 forward 函数中，我们定义了输入和输出之间的计算。接下来，我们创建了一个 net 实例，并通过 to 函数将其移至 GPU 上进行训练。

四、 模型训练

在神经网络模型搭建完成后，我们需要对其进行训练，并评估模型的性能。在 PyTorch 中，我们可以使用 nn.CrossEntropyLoss 函数来计算分类损失，并使用反向传播算法进行模型参数的更新。

下面是一个示例代码，它展示了如何使用 PyTorch 进行模型训练。

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data[0].to(device), data[1].to(device)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 100))

上述代码中，我们首先定义了损失函数和优化器。在 for 循环中，我们遍历训练数据集，将输入和标签移至 GPU 上，并使用 optimizer.zero_grad() 函数清空梯度信息；然后我们计算输出，根据损失函数计算损失并进行反向传播，最后使用 optimizer.step() 函数更新模型参数。运行过程中，我们统计并输出损失值。

五、 模型评估

在模型训练后，我们需要对其进行评估，以了解模型的性能如何。在 PyTorch 中，我们可以使用模型的 eval() 函数来将其设置为评估模式，并使用 nn.functional.softmax 函数计算对数概率值。

下面是一个示例代码，它演示了如何使用 PyTorch 进行模型评估。

correct = 0
total = 0
with torch.no_grad():
    for data in testloader:
        images, labels = data[0].to(device), data[1].to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))

上述代码中，我们使用 eval() 函数将模型设置为评估模式，并遍历测试数据集，使用 nn.functional.softmax 函数计算对数概率值，然后使用 torch.max 函数计算出最大值及其对应的位置，最后计算准确率。

总结

本篇文章从安装和配置、数据预处理、神经网络搭建、模型训练、模型评估等多个方面详细介绍了 CUDA 11.5 对应的 PyTorch。希望本篇文章可以对深度学习开发者有所帮助。