在 Ubuntu 上安装 PyTorch GPU

随着深度学习的兴起，PyTorch作为一款强大的开源机器学习库，受到了越来越多研究者和开发者的青睐。在本篇文章中，我们将在 Ubuntu 操作系统上安装支持 GPU 的 PyTorch，并展示如何进行简单的实验。同时，文中将包含一些使用代码的示例。

安装前的准备

在安装 PyTorch 之前，我们需要确认系统是否支持 NVIDIA GPU。我们可以通过以下命令检查 GPU 是否成功安装：

lspci | grep -i nvidia

若系统成功识别到了 NVIDIA GPU，那么接下来我们需要安装 NVIDIA 驱动和 CUDA Toolkit。这是使用 GPU 加速 PyTorch 所必需的。

安装 NVIDIA 驱动

首先，我们需要更新系统的包管理器：

sudo apt update
sudo apt upgrade

接下来，添加 NVIDIA 的 PPA（个人包档案）来获取最新的驱动：

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt update

然后，我们可以安装推荐的 NVIDIA 驱动：

sudo apt install nvidia-driver-<version>

请将 <version> 替换为你要安装的驱动版本，如 470。

安装 CUDA Toolkit

可以通过以下命令安装 CUDA Toolkit：

sudo apt install nvidia-cuda-toolkit

安装完成后，重启计算机以使驱动生效。

安装 PyTorch

现在我们已经设置好了 CUDA 环境，接下来可以通过 pip 安装 PyTorch。打开终端，创建一个虚拟环境是一个良好的做法：

python3 -m venv myenv
source myenv/bin/activate

然后我们用以下命令安装 PyTorch：

pip install torch torchvision torchaudio --extra-index-url

请根据你的 CUDA 版本选择合适的 <version>（例如 117，对应 CUDA 11.7）。

验证安装

为了确认 PyTorch 是否成功安装并能够使用 GPU，我们可以运行以下 Python 代码：

import torch

# 检查 GPU 是否可用
if torch.cuda.is_available():
    print(f'CUDA可用，GPU索引为: {torch.cuda.current_device()}')
    print(f'GPU名称为: {torch.cuda.get_device_name(torch.cuda.current_device())}')
else:
    print('CUDA不可用')

运行上述代码后，若一切顺利，你将会看到类似如下的输出：

CUDA可用，GPU索引为: 0
GPU名称为: NVIDIA GeForce GTX...

通过 PyTorch 进行简单的深度学习任务

安装完成后，我们可以进行简单的深度学习实验。例如，下面是一个使用 PyTorch 创建简单神经网络进行 MNIST 手写数字识别的示例：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 超参数设定
batch_size = 64
learning_rate = 0.01
num_epochs = 5

# 数据准备
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)

# 定义简单的神经网络
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)  # 输入层
        self.fc2 = nn.Linear(128, 10)        # 输出层

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 模型训练
model = SimpleNN().cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    for data, target in train_loader:
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

    print(f'Epoch {epoch + 1}, Loss: {loss.item()}')

这段代码展示了如何构建一个简单的神经网络，并使用 MNIST 数据集进行训练。

结尾

通过上述步骤，我们成功在 Ubuntu 上安装了支持 GPU 的 PyTorch，并运行了简单的示例代码。希望这篇文章能帮助你快速入门 PyTorch 的学习和使用，不论你是研究者、学生，还是开发者，都能在深度学习的旅程中找到乐趣和成就感。

journey
    title 安装 PyTorch 的旅程
    section 1. 准备工作
      检查 GPU: 5: 了解是否具备支持条件
      更新系统: 4: 更新包以安装最新驱动
    section 2. 安装
      安装 NVIDIA 驱动: 3: 安装最新驱动
      安装 CUDA: 4: 设置加速环境
      安装 PyTorch: 5: 使用 pip 安装
    section 3. 运行
      验证 GPU: 5: 确保环境正常
      进行模型训练: 4: 开始深度学习实践

相信只有实践才能不断提升技能，让我们希望你在这条机器学习的路上，走得越远越好！