如何在Ubuntu上运行PyTorch程序

在 Ubuntu 上运行 PyTorch 的完整流程

一 准备环境

• 更新系统并安装基础工具：
  • sudo apt update && sudo apt upgrade
  • sudo apt install -y python3 python3-pip build-essential
• 建议使用虚拟环境隔离依赖：
  • 使用 venv：python3 -m venv venv && source venv/bin/activate
  • 或使用 conda：conda create -n pytorch python=3.10 -y && conda activate pytorch
• 升级 pip：python -m pip install -U pip
• 可选：如需多版本 Python，可用 pyenv 管理版本后再创建虚拟环境。

二 安装 PyTorch

• 选择安装方式（CPU 或 GPU）。优先从 PyTorch 官网获取与你系统、CUDA 版本匹配的命令：https://pytorch.org/get-started/locally/
• 常见安装命令示例：
  • 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，含 cudatoolkit）：conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch
  • 使用 conda（CPU only）：conda install pytorch torchvision torchaudio cpuonly -c pytorch
• 安装完成后建议立即进行验证（见下一节）。

三 验证安装与运行第一个程序

• 验证安装（在已激活的虚拟环境中执行）：
  • python - <<‘PY’ import torch print(“PyTorch version:”, torch.version) print(“CUDA available:”, torch.cuda.is_available()) if torch.cuda.is_available(): print(“CUDA device count:”, torch.cuda.device_count()) print(“Current device:”, torch.cuda.current_device()) print(“Device name:”, torch.cuda.get_device_name(0)) PY
• 运行示例脚本 test.py：
  • nano test.py
  • 粘贴以下内容并保存：
    • import torch x = torch.rand(5, 3) y = torch.rand(5, 3) z = torch.matmul(x, y) print(z)
  • 执行：python test.py
• 启用 GPU（如果可用）的最小改动：
  • device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
  • model.to(device); x = x.to(device); y = y.to(device)。

四 GPU 环境检查与常见问题

• 检查 NVIDIA 驱动与工具链：
  • 驱动：nvidia-smi（能显示驱动版本与 GPU 列表）
  • CUDA 编译器：nvcc --version（若提示未找到，说明未安装或未加入 PATH）
• 驱动与 CUDA/cuDNN 的关系与提示：
  • 使用 pip 安装 GPU 版 PyTorch 时，官方预编译包通常已包含所需 CUDA 运行时，一般无需单独安装完整 CUDA Toolkit；但需确保 NVIDIA 驱动版本满足要求。
  • 若选择 conda 安装并指定 cudatoolkit，conda 会管理对应版本的 CUDA 运行时依赖。
  • 若需要源码编译或特定版本，才可能需要自行配置 CUDA/cuDNN 与相应环境变量（如 PATH、LD_LIBRARY_PATH）。

五 部署与进阶

• 以服务方式运行：
  • 使用 Flask/Django 将模型封装为 HTTP 服务，示例：
    • from flask import Flask, request, jsonify import torch app = Flask(name) @app.route(‘/predict’, methods=[‘POST’]) def predict(): data = request.json[‘data’] input_tensor = torch.tensor(data, dtype=torch.float32).unsqueeze(0) with torch.no_grad(): output = model(input_tensor) _, predicted = torch.max(output, 1) return jsonify({‘prediction’: predicted.item()}) if name == ‘main’: app.run(host=‘0.0.0.0’, port=5000)
• 容器化部署：
  • 示例 Dockerfile：
    • FROM python:3.10-slim WORKDIR /app COPY requirements.txt . RUN pip install -r requirements.txt COPY . . CMD [“python”, “app.py”]
  • 构建与运行：
    • docker build -t my-pytorch-app .
    • docker run -p 5000:5000 my-pytorch-app
• 生产建议：
  • 使用 Gunicorn/uWSGI + Nginx 托管 Flask 服务
  • 模型推理使用 torch.no_grad()、批处理与合适的 num_workers
  • 固定依赖版本（requirements.txt/conda env.yml），便于复现与回滚。