在 Ubuntu 上使用 PyTorch 进行深度学习的实操指南

一 环境准备与安装

• 系统更新与基础依赖
  • 更新软件源并安装基础工具：sudo apt update && sudo apt install -y python3 python3-pip python3-venv build-essential git cmake
• 创建隔离环境（推荐）
  • 使用 venv：python3 -m venv ~/pytorch_env && source ~/pytorch_env/bin/activate
  • 或使用 conda：conda create -n pytorch_env python=3.10 -y && conda activate pytorch_env
• 安装 PyTorch（选择其一）
  • CPU 版：pip install torch torchvision torchaudio
  • GPU 版（示例为 CUDA 11.8）：pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
  • 或使用 conda（GPU）：conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
• 验证安装
  • python - <<‘PY’ import torch print(“PyTorch version:”, torch.version) print(“CUDA available:”, torch.cuda.is_available()) if torch.cuda.is_available(): print(“CUDA device:”, torch.cuda.get_device_name(0)) PY
• 备注
  • 若使用 NVIDIA GPU，需确保已安装匹配的 NVIDIA 驱动、CUDA Toolkit、cuDNN，版本要与所选 PyTorch 版本兼容；安装命令以官网生成器为准更稳妥。

二 GPU 环境配置要点

• 驱动与工具链
  • 建议通过官方渠道或 Ubuntu 附加驱动安装 NVIDIA 驱动；随后安装与驱动匹配的 CUDA 与 cuDNN，并设置环境变量（如 PATH、LD_LIBRARY_PATH）以便运行时可找到相关库。
• 版本匹配与稳定性
  • 经验上 Ubuntu 22.04 + CUDA 11.8/12.x 的组合更稳；Ubuntu 24.04 对新版本框架的兼容性在部分场景下仍有问题，部署前建议先做小版本验证。
• 快速自检
  • nvidia-smi 可查看驱动与 GPU 状态；nvcc --version 检查 CUDA 编译器版本；在 Python 中 torch.cuda.is_available() 应返回 True。

三 训练一个 MNIST 卷积神经网络

• 完整可运行示例（单文件 mnist_cnn.py）

  • python - <<‘PY’ import torch, torch.nn as nn, torch.nn.functional as F from torch.utils.data import DataLoader from torchvision import datasets, transforms

    1) 设备

    device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)

    2) 数据

    transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_ds = datasets.MNIST(“./data”, train=True, download=True, transform=transform) train_loader = DataLoader(train_ds, batch_size=64, shuffle=True, num_workers=2)

    3) 模型

    class Net(nn.Module): def init(self): super().init() self.conv1 = nn.Conv2d(1, 10, 5) self.conv2 = nn.Conv2d(10, 20, 5) self.conv2_drop = nn.Dropout2d(0.25) self.fc1 = nn.Linear(320, 50) self.fc2 = nn.Linear(50, 10) def forward(self, x): x = F.relu(F.max_pool2d(self.conv1(x), 2)) x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2)) x = x.view(-1, 320) x = F.relu(self.fc1(x)) x = F.dropout(x, training=self.training) x = self.fc2(x) return F.log_softmax(x, dim=1)

    model = Net().to(device) optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5) criterion = nn.NLLLoss()

    4) 训练

    model.train() for epoch in range(2): for batch_idx, (data, target) in enumerate(train_loader, 1): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if batch_idx % 500 == 0: print(f"Epoch {epoch} [{batch_idx*len(data):>6}/{len(train_loader.dataset)}] Loss: {loss.item():.4f}")

    5) 保存与加载

    torch.save(model.state_dict(), “mnist_cnn.pth”) model2 = Net().to(device) model2.load_state_dict(torch.load(“mnist_cnn.pth”)) model2.eval() print(“Saved and reloaded model, ready for evaluation.”) PY

• 运行与扩展
  • 训练完成后可用 torch.save / torch.load 持久化模型；在验证/测试阶段使用 model.eval() 与 torch.no_grad() 提升效率与正确性。

四 常见问题与优化建议

• 环境与版本
  • 避免随意升级系统内核或驱动导致 驱动/CUDA/框架 不匹配；遇到难以排查的问题时，优先回退到已验证的 Ubuntu LTS + 稳定驱动/CUDA 组合。
• 包管理与渠道
  • 优先使用 虚拟环境 隔离依赖；GPU 版本建议选择 pip 官方 whl 或 conda + nvidia 通道 的预编译包，减少自行编译成本与不确定性。
• 性能与可复现
  • 使用 DataLoader 的 num_workers>0 加速数据读取；固定随机种子（torch.manual_seed / torch.cuda.manual_seed_all）以提升实验可复现性；必要时开启 混合精度训练 与合适的 批量大小 提升吞吐。