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-venv
• 建议始终在虚拟环境中部署，避免依赖冲突：
  • python3 -m venv venv && source venv/bin/activate
• 如需GPU支持，准备与所选PyTorch版本匹配的NVIDIA驱动；若使用pip安装带CUDA的PyTorch，包内自带所需CUDA运行时，但驱动仍需系统层面兼容。

二 安装PyTorch

• 选择安装方式（pip或conda），并严格匹配torch / torchvision / torchaudio的版本组合；GPU包需与CUDA版本一致。常见组合示例（以实际环境为准）：
  • CPU版（pip）：pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
  • GPU版（pip，CUDA 12.1）：pip3 install torch==2.1.0+cu121 torchvision==0.16.0+cu121 torchaudio==2.1.0+cu121 --index-url https://download.pytorch.org/whl/cu121
  • GPU版（pip，CUDA 12.6）：pip3 install torch==2.6.0+cu126 torchvision==0.17.0+cu126 torchaudio==2.6.0+cu126 --index-url https://download.pytorch.org/whl/cu126
  • Conda版（CPU）：conda install pytorch torchvision torchaudio cpuonly -c pytorch
  • Conda版（GPU，CUDA 11.8）：conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch -c nvidia
• 安装完成后立即验证：
  • python - <<‘PY’ import torch print(“PyTorch版本:”, torch.version) print(“CUDA可用:”, torch.cuda.is_available()) print(“GPU数量:”, torch.cuda.device_count()) PY

三 模型准备与推理服务

• 加载与推理的最小示例（请将路径与预处理替换为你的实际逻辑）：
  • python - <<‘PY’ import torch from model import MyModel # 你的模型定义 model = MyModel() model.load_state_dict(torch.load(“model.pth”, map_location=“cpu”)) # 无GPU时加载到CPU model.eval() def preprocess(x): return torch.tensor(x, dtype=torch.float32) # TODO: 按你的数据实现 def postprocess(y): return y.argmax(dim=-1).item() # TODO: 按你的任务实现 x = preprocess([[0.1]*784]) # 示例输入 with torch.no_grad(): y = model(x) print(“Prediction:”, postprocess(y)) PY
• 生产级服务化可选方案：
  • Web API：使用Flask/FastAPI封装推理接口；配合Gunicorn/uWSGI（WSGI）或Uvicorn（ASGI）部署。
  • 高性能与托管：TorchServe用于模型管理、批处理与多模型服务；ONNX Runtime用于跨平台高性能推理。

四 容器化与生产部署

• 使用Docker将应用与依赖打包，便于复现与扩缩容。示例：
  • Dockerfile
    • FROM python:3.10-slim
    • WORKDIR /app
    • COPY requirements.txt .
    • RUN pip install --no-cache-dir -r requirements.txt
    • COPY . .
    • CMD [“gunicorn”, “app:app”, “-b”, “0.0.0.0:5000”, “–timeout”, “120”]
  • requirements.txt（示例）
    • torch==2.6.0
    • torchvision==0.17.0
    • torchaudio==2.6.0
    • flask
    • gunicorn
  • 构建与运行
    • docker build -t my-pytorch-app .
    • docker run --gpus all -p 5000:5000 my-pytorch-app # 有NVIDIA GPU时启用–gpus
• 无GPU环境可移除“–gpus all”，或在requirements中使用CPU版PyTorch。

五 常见问题与排查

• 版本不匹配：确保torch / torchvision / torchaudio版本一致；GPU包与CUDA版本一致（如cu121/cu126）。
• CUDA不可用：优先检查系统NVIDIA驱动是否安装且与CUDA版本匹配；使用conda时可用cudatoolkit管理；使用pip时确认安装了带CUDA的wheel并再次运行torch.cuda.is_available()。
• 网络与安装速度：pip安装可使用国内镜像源（如清华源）加速；若遇到编译或依赖问题，先升级pip并安装常见系统依赖（build-essential、cmake、blas/lapack等）。
• 资源与稳定性：大模型推理注意显存/内存与超时设置；生产建议使用Gunicorn/Uvicorn的工作进程与超时调优，并配合日志与监控。