PyTorch 模型部署在Linux上面的详细步骤介绍

PyTorch 模型部署在Linux上面的详细步骤介绍
在 Linux 上部署 PyTorch 模型是生产环境中最常见、最成熟的方案。Linux 提供了最好的性能优化支持(如 TensorRT、ONNX Runtime)、最稳定的 Docker 容器化支持以及最完善的集群管理工具(Kubernetes)。
本文介绍了在Linux系统上部署PyTorch模型的三种主要方案。首先强调了环境准备的重要性,包括检查NVIDIA驱动、安装匹配的PyTorch版本以及验证GPU可用性。重点推荐Docker容器化部署方案,详细说明了Dockerfile编写、镜像构建和GPU启用的具体步骤。此外,还介绍了TorchServe和Triton两种高性能推理服务方案,以及传统的FastAPI+Gunicorn部署方式。文章针对每种方案都提供了具体实现步骤和命令示例,帮助开发者在Linux环境中高效部署PyTorch模型。
以下是针对 Linux 环境的详细部署指南,涵盖从环境准备到高性能推理的全流程。
前置准备:环境检查
在 Linux 上,首要任务是确保 CUDA 驱动 和 PyTorch CUDA 版本 匹配。
- 检查 NVIDIA 驱动:
nvidia-smi # 输出应显示 GPU 型号、驱动版本和 CUDA Version (这是驱动支持的最高版本)
- 安装匹配的 PyTorch:
- 前往 PyTorch 官网 获取命令。
- 关键点:确保
cudatoolkit或cuda版本小于等于nvidia-smi显示的版本。
# 示例 (根据实际版本调整) pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
- 验证 GPU 可用性:
import torch print(f"CUDA Available: {torch.cuda.is_available()}") print(f"GPU Count: {torch.cuda.device_count()}") print(f"GPU Name: {torch.cuda.get_device_name(0)}")
方案一:Docker 容器化部署 (行业标准,强烈推荐)
这是 Linux 部署的最佳实践。它解决了依赖冲突、环境隔离和可移植性问题。
步骤 1: 准备 Dockerfile
创建一个 Dockerfile,基于官方的 PyTorch 镜像(已预装 CUDA 和 cuDNN)。
# 使用官方 PyTorch 镜像 (包含 CUDA 11.8, Python 3.10)
FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
# 设置工作目录
WORKDIR /app
# 安装额外依赖 (如 FastAPI, Pillow 等)
RUN pip install --no-cache-dir fastapi uvicorn[standard] pillow python-multipart onnxruntime-gpu
# 复制模型文件和代码
# 假设你已经导出了 model.pt (见下文 TorchScript 部分)
COPY resnet18_linux.pt ./model.pt
COPY main.py ./main.py
# 暴露端口
EXPOSE 8000
# 启动命令
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
步骤 2: 构建镜像
docker build -t my-pytorch-service .
步骤 3: 运行容器 (启用 GPU)
关键参数:--gpus all (需要安装 nvidia-container-toolkit)。
docker run --gpus all -p 8000:8000 --name pt-server my-pytorch-service
如果报错 could not select device driver,请安装 NVIDIA Container Toolkit:
# Ubuntu 示例
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker
优势
- 一致性:开发环境和生产环境完全一致。
- 扩展性:轻松对接 Kubernetes (K8s) 进行自动扩缩容。
- 隔离:不污染宿主机的 Python 环境。
方案二:高性能推理服务 (TorchServe / Triton)
如果需要高并发、多模型版本管理或动态批处理(Dynamic Batching),不要自己写 Flask/FastAPI,直接使用专用推理服务器。
选项 A: TorchServe (PyTorch 官方)
-
安装:
pip install torchserve torch-model-archiver torch-workflow-archiver # 或者使用 Docker docker pull pytorch/torchserve -
打包模型 (.mar):
你需要准备model.py(定义 Handler) 和.pt文件。torch-model-archiver --model-name resnet18 \ --version 1.0 \ --model-file model.py \ --serialized-file resnet18_linux.pt \ --handler image_classifier \ --extra-files "index_to_name.json"这会生成
resnet18.mar文件。 -
启动服务:
torchserve --start --model-store model_store --models resnet18=resnet18.mar --disable_token_auth -
测试:
curl http://localhost:8080/predictions/resnet18 -T test_image.jpg
选项 B: NVIDIA Triton Inference Server (极致性能)
如果你的服务器是 NVIDIA GPU 且追求极致吞吐量,Triton 是首选。它支持 TensorRT 加速、动态批处理和并发模型执行。
- 步骤:将 PyTorch 模型导出为 ONNX 或 TensorRT Engine,然后配置
config.pbtxt,通过 Docker 运行 Triton Server。 - 优势:比原生 PyTorch 推理快 2-5 倍,显存利用率更高。
方案三:传统 Python 服务 (FastAPI + Gunicorn)
如果你不想用 Docker,或者需要在现有 Linux 服务器上快速部署。
步骤 1: 创建虚拟环境
python3 -m venv venv
source venv/bin/activate
pip install torch torchvision fastapi uvicorn gunicorn python-multipart
步骤 2: 编写 main.py (带设备自动检测)
import torch
from fastapi import FastAPI, File, UploadFile
from PIL import Image
import io
import torchvision.transforms as transforms
app = FastAPI()
# 自动选择设备:优先 CUDA,其次 CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Running on: {device}")
# 加载模型 (建议先导出为 TorchScript)
model = torch.jit.load("resnet18_linux.pt")
model.to(device)
model.eval()
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
img = Image.open(io.BytesIO(await file.read())).convert("RGB")
input_tensor = transform(img).unsqueeze(0).to(device)
with torch.no_grad():
output = model(input_tensor)
return {"class_id": int(output.argmax(dim=1).item())}
步骤 3: 使用 Gunicorn 启动 (生产级 WSGI/ASGI 服务器)
不要直接用 python main.py 或 uvicorn 单进程运行。使用 Gunicorn 管理多 Worker。
# -w: worker 数量 (通常设为 CPU 核心数)
# -k: worker 类 (uvicorn.workers.UvicornWorker 用于异步 ASGI)
# --timeout: 防止长推理请求被杀掉
gunicorn -w 4 -k uvicorn.workers.UvicornWorker --timeout 120 --bind 0.0.0.0:8000 main:app
步骤 4: 设置为 Systemd 服务 (开机自启)
创建 /etc/systemd/system/pytorch-api.service:
[Unit]
Description=PyTorch FastAPI Service
After=network.target
[Service]
User=ubuntu
Group=ubuntu
WorkingDirectory=/home/ubuntu/my_project
Environment="PATH=/home/ubuntu/my_project/venv/bin"
ExecStart=/home/ubuntu/my_project/venv/bin/gunicorn -w 4 -k uvicorn.workers.UvicornWorker --bind 0.0.0.0:8000 main:app
Restart=always
RestartSec=10
[Install]
WantedBy=multi-user.target
启动服务:
sudo systemctl daemon-reload
sudo systemctl start pytorch-api
sudo systemctl enable pytorch-api # 开机自启
sudo systemctl status pytorch-api
方案四:C++ 高性能部署 (LibTorch)
适用于对延迟极其敏感、需要嵌入到现有 C++ 后端服务的场景。
- 下载 LibTorch (Linux CUDA 版):
从 PyTorch 官网下载LibTorch (CUDA 11.x)的.zip包。wget https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.1.0%2Bcu118.zip unzip libtorch*.zip
- 编写 CMakeLists.txt:
cmake_minimum_required(VERSION 3.0 FATAL_ERROR) project(custom_ops) find_package(Torch REQUIRED) add_executable(example-app example-app.cpp) target_link_libraries(example-app "${TORCH_LIBRARIES}") set_property(TARGET example-app PROPERTY CXX_STANDARD 17)
- 编译与运行:
注意:运行时需要将 libtorch 的mkdir build && cd build cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch .. make ./example-applib目录加入LD_LIBRARY_PATH。export LD_LIBRARY_PATH=/path/to/libtorch/lib:$LD_LIBRARY_PATH
Linux 部署特别注意事项与优化
- 权限与用户:
- 不要用
root运行推理服务。创建专用用户(如inference),并赋予其对模型文件和日志目录的读写权限。 - GPU 访问权限:默认情况下,普通用户可能无法访问
/dev/nvidia*。需将用户加入video或render组,或检查 Docker 配置。
- 不要用
- 性能优化技巧:
- 线程数控制:Linux 上 PyTorch 默认可能占用所有 CPU 核心用于数据加载,导致与推理争抢资源。
torch.set_num_threads(1) # 在推理线程中设置 - 混合精度推理 (AMP):如果显卡支持 Tensor Core (Volta 架构及以后),使用
torch.autocast可以提速并减少显存。with torch.autocast(device_type="cuda", dtype=torch.float16): output = model(input_tensor) - NUMA 绑定:在多路 CPU 服务器上,使用
numactl绑定进程到特定的 CPU 核和内存节点,减少跨节点通信延迟。numactl --cpunodebind=0 --membind=0 python main.py
- 线程数控制:Linux 上 PyTorch 默认可能占用所有 CPU 核心用于数据加载,导致与推理争抢资源。
- 监控与日志:
- 使用
nvtop(类似 htop 的 GPU 监控工具) 实时监控显存和利用率。 - 集成 Prometheus + Grafana,通过
prometheus-fastapi-instrumentator暴露 QPS、延迟、错误率等指标。
- 使用
- 模型热更新:
- 在生产环境中,避免重启服务来更新模型。可以在代码中实现监听模型文件变化,或使用 TorchServe/Triton 的版本管理功能 (
model_version) 动态加载新版本。
- 在生产环境中,避免重启服务来更新模型。可以在代码中实现监听模型文件变化,或使用 TorchServe/Triton 的版本管理功能 (
总结推荐
| 场景 | 推荐方案 | 理由 |
|---|---|---|
| 通用生产环境 | 方案一 (Docker + FastAPI) | 标准化、易维护、易扩展,兼容 K8s。 |
| 高并发/企业级 | 方案二 (Triton / TorchServe) | 支持动态批处理、多模型、TensorRT 加速,吞吐量最高。 |
| 老旧服务器/无 Docker | 方案三 (Systemd + Gunicorn) | 直接利用系统服务管理,稳定可靠。 |
| 超低延迟/嵌入式 | 方案四 (C++ LibTorch) | 去除 Python 开销,极致性能,体积小。 |
对于大多数 Linux 部署任务,Docker 容器化 + FastAPI 是性价比最高、最稳妥的选择;若对性能有极致要求,则迁移至 Triton Inference Server。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐


所有评论(0)