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 torch print(“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 torch from 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, jsonify import torch, torch.nn as nn app = 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-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 --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）。