在这篇文章中,我们将讨论如何调优 TorchServe 以便在生产环境中部署模型。机器学习项目生命周期中最大的挑战之一是在生产环境中部署模型。这需要一个可靠的服务解决方案,并辅以满足 MLOps 需求的配套方案。一个强大的服务解决方案需要提供对多模型服务、模型版本控制、指标记录、监控以及应对峰值流量的扩展能力的支持。在本文中,我们将概述 TorchServe,并探讨如何针对生产用例调优其性能。我们还将讨论来自 Meta 的 Animated Drawings(动画绘图)应用,该应用可以将人类素描转化为动画,并展示如何通过 TorchServe 处理峰值流量。Animated Drawings 的工作流如下。

https://ai.facebook.com/blog/using-ai-to-bring-childrens-drawings-to-life
许多 AI 系统和工具旨在处理逼真的人类图像,而儿童绘画因其通常以抽象、天马行空的方式构成,增加了一层复杂性和不可预测性。这类形态和风格上的差异甚至会让那些擅长在真实照片和图画中识别物体的尖端 AI 系统感到困惑。Meta AI 的研究人员正在努力克服这一挑战,使 AI 系统能够更好地识别儿童以各种迥异方式创作的人类素描。这篇精彩的博客文章提供了关于 Animated Drawings 及其采用方法的更多详细信息。
TorchServe

图 1. TorchServe 性能调优的整体流程
模型训练完成后,需要将其集成到更大的系统中才能形成完整的应用程序,我们使用“模型服务”(Model Serving)这一术语来指代这种集成。本质上,模型服务就是让您训练好的模型可用于运行推理及后续使用。
TorchServe 是 PyTorch 首选的生产环境模型服务解决方案。它是一款高性能且可扩展的工具,能够将您的模型封装在 HTTP 或 HTTPS API 中。它拥有一个用 Java 实现的前端,负责从为服务模型分配工作线程到处理客户端与服务器之间的连接等多种任务。TorchServe 拥有一个 Python 后端,负责处理推理服务。
TorchServe 支持多模型服务、用于 AB 测试的模型版本控制、动态批处理、日志记录和指标监控。它提供了四个用于 推理 (Inference)、解释 (Explanations)、管理 (Management) 和 指标 (Metrics) 的 API。
推理 API 默认监听 8080 端口,可通过 localhost 访问,这可以在 TorchServe 配置中进行设置,从而实现从模型获取预测结果。
解释 API 在底层使用 Captum 来提供所服务模型的解释,同样监听 8080 端口。
管理 API 允许注册、取消注册和描述模型。它还允许用户扩展或缩减服务模型的 worker(工作线程)数量。
指标 API 默认监听 8082 端口,使我们能够监控所服务的模型。
TorchServe 通过支持 批量推理 和多个服务于模型的 worker,让您可以扩展模型服务并处理峰值流量。扩展可以通过 管理 API 和 配置文件设置来完成。此外,指标 API 有助于通过默认和自定义指标来监控模型服务。
其他高级设置(例如接收请求的队列长度、输入批处理的最大等待时间以及许多其他属性)均可通过在 TorchServe 启动时传入的 配置文件进行配置。
使用 TorchServe 服务模型的步骤
- 安装 TorchServe、模型归档工具及其依赖项。
- 选择一个适合您任务的默认处理程序(如图像分类等),或者编写一个 自定义处理程序。
- 使用 TorchArchive 将您的模型工件(训练好的模型检查点以及加载和运行模型所需的所有其他文件)和处理程序打包成“.mar”文件,并将其放入模型存储中。
- 启动模型服务.
- 运行推理。我们将在下文中详细讨论模型处理程序和指标。
模型处理程序
TorchServe 在后端使用处理程序(handler)来加载模型、预处理接收到的数据、运行推理并对响应进行后处理。TorchServe 中的处理程序是一个 Python 脚本,其中包含了所有的模型初始化、预处理、推理和后处理逻辑。
TorchServe 为图像分类、分割、目标检测和文本分类等多种应用提供了开箱即用的处理程序。如果您的用例不适用于默认处理程序,它也支持自定义处理程序。
自定义处理程序提供了极大的灵活性,这使 TorchServe 有潜力成为一种 多框架 服务工具。自定义处理程序允许您定义自定义逻辑来初始化模型,该逻辑也可用于加载来自其他框架(如 ONNX)的模型。
TorchServe 处理程序 由四个主要 函数 组成:initialize、preprocess、inference 和 postprocess,每个函数都返回一个列表。下方的代码片段展示了一个自定义处理程序的示例。自定义处理程序继承 自 TorchServe 中的 BaseHandler,并可以 覆盖 任何 主要函数。这是用于加载 Detectron2 模型进行图形检测的处理程序示例,该模型已导出为 TorchScript,并使用 model.half() 以 FP16 运行推理,详情见本文的另一个 章节。
class MyModelHandler(BaseHandler):
def initialize(self, context):
self.manifest = ctx.manifest
properties = ctx.system_properties
model_dir = properties.get("model_dir")
serialized_file = self.manifest["model"]["serializedFile"]
model_pt_path = os.path.join(model_dir, serialized_file)
self.device = torch.device(
"cuda:" + str(properties.get("gpu_id"))
if torch.cuda.is_available() and properties.get("gpu_id") is not None
else "cpu"
)
self.model = torch.jit.load(model_pt_path, map_location=self.device)
self.model = self.model.half()
def preprocess(self, data):
inputs = []
for request in batch:
request_body = request.get("body")
input_ = io.BytesIO(request_body)
image = cv2.imdecode(np.fromstring(input_.read(), np.uint8), 1)
input = torch.Tensor(image).permute(2, 0, 1)
input = input.to(self.device)
input = input.half()
inputs.append({"image": input})
return inputs
def inference(self,inputs):
predictions = self.model(**inputs)
return predictions
def postprocess(self, output):
responses = []
for inference_output in inference_outputs:
responses_json = {
'classes': inference_output['pred_classes'].tolist(),
'scores': inference_output['scores'].tolist(),
"boxes": inference_output['pred_boxes'].tolist()
}
responses.append(json.dumps(responses_json))
return responses
指标
生产环境中模型服务的关键组成部分是监控它们的能力。TorchServe 定期 收集 系统级 指标,并 允许 添加 自定义指标。
系统级指标 包括主机上的 CPU 利用率、可用及已用磁盘空间和内存,以及不同响应代码(例如 200-300、400-500 和 500 以上)的请求数量。自定义指标 可以如 此处 所述添加到指标中。TorchServe 将这两组指标记录到不同的日志文件中。指标默认收集位置为:
- 系统指标 – log_directory/ts_metrics.log
- 自定义指标 – log_directory/model_metrics.log
如前所述,TorchServe 还提供了 指标 API,该 API 默认监听 8082 端口,使用户能够查询和监控收集到的指标。默认指标端点返回 Prometheus 格式的指标。您可以使用 curl 请求查询指标,或者将 Prometheus 服务器 指向该端点,并使用 Grafana 进行仪表板展示。
在服务模型时,您可以按照以下方式使用 curl 请求查询指标:
curl http://127.0.0.1:8082/metrics
如果您打算导出记录的指标,请参考这个使用 mtail 将指标导出到 Prometheus 的 示例。在仪表板中跟踪这些指标,可以帮助您监控在离线基准测试中可能零星出现或难以发现的性能下降问题。
生产环境模型性能调优的考虑因素
图 1 中建议的工作流是关于如何使用 TorchServe 处理生产环境模型部署的总体思路。
在许多情况下,生产环境的模型服务是 基于 吞吐量 或 延迟 服务等级协议 (SLA) 进行 优化 的。通常,实时应用 更关注 延迟,而 离线应用 可能更看重更高的 吞吐量。
有许多主要因素会影响生产环境模型服务的性能。虽然我们在此重点关注使用 TorchServe 服务 PyTorch 模型,但大多数因素也适用于其他框架的模型。
- 模型优化:这是将模型部署到生产环境的前置步骤。这是一个非常广泛的话题,我们将在未来的系列博客中深入探讨。这包括量化、修剪以减小模型大小、在 PyTorch 中使用中间表示(如 TorchScript)、算子融合等技术。目前,torchprep 作为一个 CLI 工具提供了其中的许多技术。
- 批量推理: 指将多个输入馈送到模型中,虽然这在训练期间至关重要,但在推理阶段管理成本也非常有用。硬件加速器针对并行性进行了优化,批处理有助于饱和计算能力,通常会导致更高的吞吐量。推理的主要区别在于您不能等待太久来从客户端填充批次,这就是所谓的动态批处理。
- Worker(工作线程)数量: TorchServe 使用 worker 来服务模型。TorchServe 的 worker 是持有模型权重副本用于运行推理的 Python 进程。worker 过少意味着您无法获得足够的并行性,但 worker 过多可能会导致争用并降低端到端性能。
- 硬件: 根据模型、应用以及延迟和吞吐量预算选择合适的硬件。这可以是 TorchServe 支持的 硬件之一:CPU、GPU、AWS Inferentia。一些硬件配置旨在实现一流的性能,而另一些则更适合具有成本效益的推理。根据我们的实验发现,GPU 在较大的批次大小下表现最佳,而合适的 CPU 和 AWS Inferentia 在较小的批次大小和低延迟场景下更具成本效益。
TorchServe 性能调优最佳实践
为了在使用 TorchServe 服务时获得最佳模型性能,我们在此分享一些最佳实践。TorchServe 提供了一个 基准测试 套件,能提供有用的见解,帮助您就下述不同选择做出明智的决定。
- 第一步是 优化您的模型,PyTorch 模型优化 教程。模型优化 的选择也与所选 硬件 紧密相关。我们将在另一篇博客文章中更详细地讨论它。
- 决定 模型部署的 硬件 可能与延迟、吞吐量预算以及单次推理成本密切相关。根据模型和应用的大小,情况各不相同。例如,对于计算机视觉模型,历史上在 CPU 上运行的成本是不可接受的。然而,随着 TorchServe 最近添加了诸如 IPEX 等优化方案,这变得更加经济高效,您可以在这份调查研究 案例研究 中了解更多信息。
- TorchServe 中的 worker 是提供并行性的 Python 进程,设置 worker 数量应谨慎。默认情况下,TorchServe 启动的 worker 数量等于主机上的 VCPU 或可用 GPU 数量,这会显著增加 TorchServe 的启动时间。TorchServe 暴露了一个 配置属性 来设置 worker 数量。为了通过 多 worker 提供 高效的并行性 并避免它们争用资源,作为基准,我们 建议 在 CPU 和 GPU 上采用以下设置:CPU:在处理程序中设置
torch.set_num_threads(1),然后将 worker 数量设置为物理核心数 / 2。但最佳线程配置可以通过利用 Intel CPU 启动脚本来实现。GPU:可用 GPU 的数量可以通过 config.properties 中的 number_gpus 进行设置。TorchServe 使用轮询(round robin)方式将 worker 分配给 GPU。我们建议按照以下方式设置 worker 数量:Worker 数量 = (可用 GPU 数量) / (独立模型数量)。请注意,Ampere 之前的 GPU 无法通过 Multi-Instance GPU (MIG) 提供资源隔离。 - 批次大小 (Batch size) 直接影响延迟和吞吐量。为了更好地利用计算资源,需要增加批次大小。然而,延迟和吞吐量之间存在权衡。更大的批次大小 可以 增加吞吐量,但也会导致更高的延迟。批次大小可以通过两种方式在 TorchServe 中设置:通过 config.properties 中的 模型配置,或者在使用 管理 API 注册模型时设置。
在下一节中,我们将使用 TorchServe 基准测试套件来确定模型优化、硬件、worker 和批次大小的最佳组合。
Animated Drawings 性能调优
要使用 TorchServe 基准测试套件,首先我们需要一个归档文件(如上所述的“.mar”文件),其中包含模型、处理程序以及加载和运行推理所需的所有其他工件。Animated Drawings 使用 Detectron2 的 Mask-RCNN 实现作为目标检测模型。
如何运行基准测试套件
TorchServe 中的 自动化基准测试套件 允许您在不同设置(包括批次大小和 worker 数量)下对多个模型进行基准测试,并最终为您生成报告。入门步骤如下:
git clone https://github.com/pytorch/serve.git
cd serve/benchmarks
pip install -r requirements-ab.txt
apt-get install apache2-utils
模型级别的设置可以在 YAML 文件中进行配置,类似如下:
Model_name:
eager_mode:
benchmark_engine: "ab"
url: "Path to .mar file"
workers:
- 1
- 4
batch_delay: 100
batch_size:
- 1
- 2
- 4
- 8
requests: 10000
concurrency: 10
input: "Path to model input"
backend_profiling: False
exec_env: "local"
processors:
- "cpu"
- "gpus": "all"
此 YAML 文件将在 benchmark_config_template.yaml 文件中被引用,该文件包含用于生成报告的其他设置,它还可以选择与 AWS CloudWatch 集成以记录日志。
python benchmarks/auto_benchmark.py --input benchmark_config_template.yaml
运行 基准测试,结果将写入“_ /tmp/benchmark/ab_report.csv _”中的“csv”文件,完整报告位于“/tmp/ts_benchmark/report.md”。报告将包含 TorchServe 平均延迟、模型 P99 延迟、吞吐量、并发数、请求数量、处理程序时间以及其他一些指标。我们重点关注用于调优性能的几个关键指标:并发数、模型 P99 延迟、吞吐量。我们特别结合 批次大小、所使用的 设备、worker 数量 以及是否进行了 模型优化 来观察这些数值。
该模型的 延迟 SLA 已设定为 100 毫秒,这是一个实时应用,正如我们之前讨论的那样,延迟是首要考虑因素,而 吞吐量 在 不违反延迟 SLA 的前提下理想应尽可能高。
通过搜索不同的批次大小 (1-32)、worker 数量 (1-16) 和设备 (CPU, GPU),我们进行了一系列实验,下表总结了其中的最优配置。
| 设备 | 并发数 | 请求数 | Worker数 | 批次大小 | 载荷/图像 | 优化 | 吞吐量 | 延迟 P99 |
| CPU | 10 | 1000 | 1 | 1 | 小 | 不适用 | 3.45 | 305.3 毫秒 |
| CPU | 1 | 1000 | 1 | 1 | 小 | 不适用 | 3.45 | 291.8 毫秒 |
| GPU | 10 | 1000 | 1 | 1 | 小 | 不适用 | 41.05 | 25.48 毫秒 |
| GPU | 1 | 1000 | 1 | 1 | 小 | 不适用 | 42.21 | 23.6 毫秒 |
| GPU | 10 | 1000 | 1 | 4 | 小 | 不适用 | 54.78 | 73.62 毫秒 |
| GPU | 10 | 1000 | 1 | 4 | 小 | model.half() | 78.62 | 50.69 毫秒 |
| GPU | 10 | 1000 | 1 | 8 | 小 | model.half() | 85.29 | 94.4 毫秒 |
该模型在 CPU 上,无论在批次大小、并发数和 worker 数量方面进行何种尝试,其延迟都无法满足 SLA,实际上高出约 13 倍。
将模型服务 迁移到 GPU,可立即 将延迟 改善 约 13 倍,从 305 毫秒降低到 23.6 毫秒。
我们能为模型做的最 简单优化 之一是将其精度降低为 fp16,这只需一行代码 (model.half()),且能将 模型 P99 延迟 降低 32%,并将吞吐量增加近乎相同的幅度。
还有其他优化手段,例如将模型 TorchScript 化并使用 optimize_for_inference,或包括 ONNX 或 TensorRT 运行时优化(利用激进的融合技术)在内的其他技巧,均超出了本文的讨论范围。我们将在另一篇博客中讨论模型优化。
我们发现,无论在 CPU 还是 GPU 上,在本例中 设置 worker 数量=1 效果最好。
- 将模型迁移到 GPU,并使用 worker 数量 = 1 和 批次大小 = 1,与 CPU 相比,吞吐量增加了约 12 倍,延迟降低了约 13 倍。
- 将模型迁移到 GPU,使用 model.half()、worker 数量 = 1 和 批次大小 = 8,在 吞吐量 和可容忍延迟方面获得了 最佳 结果。吞吐量 相比 CPU 增加了约 25 倍,且延迟仍满足 SLA (94.4ms)。
注意:如果您正在运行基准测试套件,请确保设置了适当的 batch_delay,并将请求的并发数设置为与您的批次大小成比例的数字。此处的并发数是指发送到服务器的并发请求数量。
结论
在本文中,我们讨论了 TorchServe 为优化生产环境性能所提供的考量因素和配置选项。我们介绍了 TorchServe 基准测试套件,作为一种调优性能并获取模型优化、硬件选择和成本总体见解的手段。我们以使用 Detectron2 Mask-RCNN 模型的 Animated Drawings 应用为例,展示了如何通过基准测试套件进行性能调优。
有关 TorchServe 性能调优的更多详细信息,请参考我们的文档 此处。如有任何进一步的问题和反馈,也欢迎在 TorchServe 仓库 中提交工单。
致谢
我们要感谢 Somya Jain (Meta) 和 Christopher Gustave (Meta) 在本博客撰写过程中提供的巨大支持和指导,以及对 Sketch Animator 工作流的见解。此外,特别感谢来自 AWS 的 Li Ning,他通过自动化基准测试套件,为使 TorchServe 上的性能调优变得更加简单做出了巨大努力。