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 版本或参考下表示例命令。

方式	场景	安装命令示例
pip	CPU	pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip	CUDA 11.8	pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
conda	CPU	conda install pytorch torchvision torchaudio cpuonly -c pytorch
conda	CUDA 11.8	conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia

说明：请将命令中的 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）。

0 赞

0 踩