如何用Ubuntu运行PyTorch模型

在 Ubuntu 上运行 PyTorch 模型的完整流程

一环境准备

更新系统并安装基础工具：
- sudo apt update && sudo apt upgrade
- sudo apt install -y python3 python3-pip
建议使用虚拟环境隔离依赖：
- python3 -m venv venv
- source venv/bin/activate
验证基础环境：
- python3 -c “import sys; print(sys.version)”
- pip3 --version

二安装 PyTorch

CPU 版本（通用、零配置）：
- pip install torch torchvision torchaudio
GPU 版本（需 NVIDIA 显卡与对应 CUDA 驱动/工具链）：
- 先确认驱动与 CUDA：nvidia-smi（右上显示 Supported/Runtime CUDA）
- 使用 pip 安装与 CUDA 匹配的 PyTorch（示例为 CUDA 11.3 的额外索引方式）：
  - pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
- 若使用 conda，也可用：conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia（版本请按你的 CUDA 选择）
验证安装与设备可用性：
- python3 - <<‘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

三准备与加载模型

常见保存与加载方式：
- 仅权重：torch.save(model.state_dict(), “model.pth”)；加载：model.load_state_dict(torch.load(“model.pth”))
- 整个模型：torch.save(model, “model.pt”)；加载：model = torch.load(“model.pt”)
推理最小示例（请将模型类或权重路径替换为你的实际内容）：
- python3 - <<‘PY’ import torch from torchvision import models # 示例：使用 torchvision 模型
  
  方式A：使用预训练权重
  
  model = models.resnet18(pretrained=True) model.eval()
  
  方式B：加载本地权重（示例）
  
  from your_model import MyModel # 你的模型定义
  
  model = MyModel()
  
  model.load_state_dict(torch.load(“model.pth”))
  
  model.eval()
  
  随机输入模拟推理（替换为你的数据预处理）
  
  dummy = torch.randn(1, 3, 224, 224) # B, C, H, W with torch.no_grad(): out = model(dummy) print(“Output shape:”, out.shape) PY
注意：推理前务必调用 model.eval()，并使用 torch.no_grad() 关闭梯度计算以提升性能与稳定性。

四本地服务化与容器化

Flask/FastAPI 快速 API（示例要点）：
- 安装：pip install flask
- 代码示例（保存为 app.py）：
  - from flask import Flask, request, jsonify import torch, torch.nn as nn app = Flask(name) class Net(nn.Linear): pass # 替换为你的模型 model = Net(10, 1); model.load_state_dict(torch.load(“model.pth”)); model.eval() @app.route(“/predict”, methods=[“POST”]) def predict(): x = torch.tensor([[float(v) for v in request.json[“x”]]], dtype=torch.float32) with torch.no_grad(): y = model(x).item() return jsonify({“y”: y}) if name == “main”: app.run(host=“0.0.0.0”, port=5000)
- 运行：python app.py
Docker 容器化（示例 Dockerfile）：
- Dockerfile
  - FROM python:3.10-slim WORKDIR /app COPY requirements.txt . RUN pip install -r requirements.txt COPY . . CMD [“python”, “app.py”]
- requirements.txt
  - torch torchvision torchaudio
  - flask
- 构建与运行：
  - docker build -t pytorch-app .
  - docker run --gpus all -p 5000:5000 pytorch-app（GPU 需安装 nvidia-container-toolkit）
生产级可选方案：TorchServe、ONNX Runtime（适合更高吞吐与跨平台部署）。

五常见问题与排查

GPU 不可用但已安装 CUDA：
- 检查驱动：nvidia-smi；确保 PyTorch 与驱动/CUDA 版本匹配
- 容器内使用 GPU：安装 nvidia-container-toolkit 并以 --gpus all 启动
版本不匹配导致安装失败或运行异常：
- 使用 conda 安装时同时指定 cudatoolkit 与渠道（pytorch、nvidia）
- 使用 pip 时选择带 cu 后缀的 wheel，或参考 PyTorch 官网命令生成器
内存不足 OOM：
- 减小 batch size，或使用 torch.cuda.empty_cache()
- 使用半精度/混合精度（如 torch.cuda.amp）降低显存占用
预处理/后处理不一致：
- 保证训练与推理的归一化、尺寸、填充、tokenization 等完全一致
生产建议：
- 固定依赖版本（requirements.txt/conda env.yaml）
- 加入日志、监控与异常捕获，必要时使用模型服务框架（如 TorchServe）

方式A：使用预训练权重