在 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），便于复现与回滚。
