pytorch与CUDA版本匹配(基于conda管理系统)以及GPU和CPU训练模型速度对比

Conda 是一个开源的软件包管理系统和环境管理系统，它可以快速安装、运行和更新软件包及其依赖项。Anaconda 是一个包含 Conda 的数据科学平台发行版。它预装了许多常用的数据科学包和工具，适合初学者或希望快速搭建数据科学环境的用户。

anaconda安装可参考这篇文章。

例如我的cuda版本是12.1，则可在终端运行下面这条指令安装pytorch全家桶。

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

 安装完成后运行下面程序即可判断安装是否正确。

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)  # 这应该输出 'cuda' 如果 GPU 可用

这里提供一个小型的神经网络程序，用来观察CPU和GPU训练模型速度的快慢。

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

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

# 数据加载
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = datasets.CIFAR10(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

# 卷积神经网络模型
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(64 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 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, 64 * 8 * 8)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 训练函数
def train(model, device, train_loader, optimizer, criterion, num_epochs):
    model.train()
    start_time = time.time()
    for epoch in range(num_epochs):
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
    end_time = time.time()
    return end_time - start_time

# 在CPU上训练
device_cpu = torch.device("cpu")
model_cpu = ConvNet().to(device_cpu)
optimizer_cpu = optim.SGD(model_cpu.parameters(), lr=learning_rate)
criterion_cpu = nn.CrossEntropyLoss()

cpu_time = train(model_cpu, device_cpu, train_loader, optimizer_cpu, criterion_cpu, num_epochs)
print(f"CPU training time: {cpu_time:.2f} seconds")

# 在GPU上训练
if torch.cuda.is_available():
    device_gpu = torch.device("cuda")
    model_gpu = ConvNet().to(device_gpu)
    optimizer_gpu = optim.SGD(model_gpu.parameters(), lr=learning_rate)
    criterion_gpu = nn.CrossEntropyLoss()

    gpu_time = train(model_gpu, device_gpu, train_loader, optimizer_gpu, criterion_gpu, num_epochs)
    print(f"GPU training time: {gpu_time:.2f} seconds")
else:
    print("GPU is not available.")

在我的电脑上运行结果如下： 

可以看出GPU可以显著提高模型的训练速度。