PyTorch于Ubuntu上的部署方法

Ubuntu上部署PyTorch的完整流程

一 环境准备与依赖

• 更新系统并安装基础工具：
  • sudo apt update && sudo apt upgrade -y
  • sudo apt install -y build-essential cmake git wget unzip yasm pkg-config libopenblas-dev liblapack-dev libjpeg-dev libpng-dev python3 python3-pip
• 使用虚拟环境隔离依赖（推荐）：
  • python3 -m venv venv && source venv/bin/activate
• GPU环境需准备与驱动匹配的 NVIDIA 驱动、CUDA、cuDNN，版本需与所选 PyTorch 版本兼容；无 NVIDIA 显卡可直接使用 CPU 版本。

二 安装PyTorch

• 选择安装方式（pip 或 conda），并严格匹配 CUDA 版本；若不确定，优先使用 CPU 版本或参考下表示例命令。

• 说明：请将命令中的 CUDA 版本替换为与你驱动/硬件匹配的版本；如需其他版本，请到 PyTorch 官网获取对应命令。

三 验证安装与环境检查

• 基础验证：
  • python - <<‘PY’import torchprint(“PyTorch version:”, torch.version)print(“CUDA available:”, torch.cuda.is_available())print(“GPU count:”, torch.cuda.device_count())if torch.cuda.is_available():print(“Current device:”, torch.cuda.get_device_name(0))PY
• 预期：输出版本号；GPU 环境应显示 CUDA available: True、设备数量与名称；CPU 环境显示 False。

四 部署方式与示例

• 本地脚本推理
  • 保存模型与权重（如 model.pth），编写推理脚本：
    • python - <<‘PY’import torchfrom model import MyModel# 你的模型定义model = MyModel()model.load_state_dict(torch.load(“model.pth”, map_location=“cpu”))# GPU 可用时改为 “cuda”model.eval()x = torch.randn(1, 3, 224, 224)# 示例输入with torch.no_grad():y = model(x)print(y.shape)PY
• Web 服务 API（Flask 示例）
  • 安装：pip install flask
  • 服务代码：
    • python - <<‘PY’from flask import Flask, request, jsonifyimport torch, torch.nn as nnapp = Flask(name)class Net(nn.Module):def init(self): super().init(); self.l = nn.Linear(10, 1)def forward(self, x): return self.l(x)model = Net(); model.load_state_dict(torch.load(“model.pth”, map_location=“cpu”)); model.eval()@app.route(“/predict”, methods=[“POST”])def predict():data = torch.tensor([[float(x) for x in request.json[“x”]]], dtype=torch.float32)with torch.no_grad(): pred = model(data).item()return jsonify({“prediction”: pred})if name == “main”:app.run(host=“0.0.0.0”, port=5000)PY
  • 运行：python app.py；访问 http://:5000/predict
• 容器化部署（Docker）
  • Dockerfile：
    • FROM python:3.9-slimWORKDIR /appCOPY requirements.txt .RUN pip install -r requirements.txtCOPY . .CMD [“python”, “app.py”]
  • 构建与运行：
    • docker build -t my-pytorch-app .
    • docker run --gpus all -p 5000:5000 my-pytorch-app# 仅 GPU 环境加 --gpus all
• 高性能与服务化
  • 大规模或低延迟场景可用 TorchServe 部署；跨框架与多平台推理可用 ONNX Runtime。

五 常见问题与排障

• 版本不匹配
  • 现象：torch.cuda.is_available() 为 False 或报错；处理：确认 驱动、CUDA、cuDNN 与 PyTorch 版本匹配，必要时重装对应版本。
• 网络安装慢或失败
  • 处理：临时使用镜像源（如清华源）或更换网络；pip 安装时添加 -i 参数。
• 内存不足
  • 处理：关闭不必要进程、增大交换空间，或使用更轻量的模型/批大小。
• GPU 不可见
  • 处理：检查 nvidia-smi 输出、驱动与容器 GPU 配置（如 Docker 需 --gpus all）。