pytorch/torchserve-nightly Docker 镜像 - 轩辕镜像

PyTorch Serve 镜像文档

概述

PyTorch Serve 是由 PyTorch 官方开发的模型服务框架，旨在简化 PyTorch 模型的部署流程，提供高性能、可扩展的推理服务能力。该框架支持从模型打包、服务启动到推理请求处理的全流程管理，适用于企业级 AI 服务部署场景。

核心特性

• 多模型格式支持：兼容 TorchScript、ONNX 等主流模型格式，支持自定义模型序列化方案
• 双接口协议：提供 REST API（推理、管理）和 gRPC API（高性能推理），满足不同场景需求
• 动态批处理：自动优化推理请求批处理策略，平衡吞吐量与延迟
• 模型版本控制：支持多版本模型并行部署，便于 A/B 测试与灰度发布
• 监控集成：内置 Prometheus 指标导出，支持推理性能、资源占用等数据监控
• 自定义处理流程：允许通过 Python 脚本扩展请求预处理/后处理逻辑
• 多资源适配：支持 CPU、GPU 部署，可配置资源分配策略

使用场景

• 企业级 AI 推理服务部署（如推荐系统、图像识别 API）
• 实时推理场景（需低延迟、高并发处理）
• 多模型管理与版本控制（如模型迭代测试）
• 边缘计算与云端协同推理服务

Docker 部署示例

前提条件

• 安装 Docker（若使用 GPU，需配置 nvidia-docker）
• 模型文件（如 TorchScript 格式 .pt 文件）

步骤 1：拉取镜像

bash
# CPU 版本
docker pull pytorch/torchserve:latest
# GPU 版本（需匹配 CUDA 版本）
docker pull pytorch/torchserve:latest-gpu

步骤 2：打包模型

使用 model-archiver 工具将模型打包为 .mar 文件（镜像内置该工具）：

bash
# 假设模型文件位于当前目录 ./model
# 打包命令（自定义模型名称、版本、处理器）
docker run --rm -v $(pwd)/model:/workspace pytorch/torchserve:latest torch-model-archiver \
  --model-name my_model \
  --version 1.0 \
  --model-file /workspace/model.py \
  --serialized-file /workspace/model.pt \
  --handler image_classifier \
  --export-path /workspace/model-store

生成的 .mar 文件将保存至 ./model/model-store 目录。

步骤 3：启动服务容器

bash
# 映射推理端口（8080）、管理端口（8081）及模型存储目录
docker run -d \
  --name torchserve-service \
  -p 8080:8080 \
  -p 8081:8081 \
  -v $(pwd)/model/model-store:/home/model-server/model-store \
  pytorch/torchserve:latest \
  torchserve --start --model-store model-store --models my_model=my_model.mar

步骤 4：测试推理

通过 REST API 发送推理请求：

bash
# 发送图像文件进行分类推理
curl http://localhost:8080/predictions/my_model -T test_image.jpg

注意事项

• 模型打包时需指定正确的 handler（如 image_classifier、object_detector），或自定义 handler 脚本
• 端口映射：8080（推理端口）、8081（管理端口，用于模型加载/卸载）
• GPU 部署需确保主机已安装 NVIDIA 驱动及 nvidia-container-toolkit
• 可通过 --ts-config 参数加载配置文件，优化批处理大小、线程数等性能参数