跳转到主要内容
博客

优化基于 LibTorch 的推理引擎内存使用和线程池

概述

在这篇博文中，我们将展示如何优化基于 LibTorch 的推理引擎，通过减少内存使用和优化线程池策略来最大化吞吐量。我们将这些优化应用于音频数据（例如音乐和语音识别或声学指纹识别）的模式识别引擎。这篇博文中讨论的优化可以将内存使用量减少 50%，并将推理的端到端延迟减少 37.5%。这些优化适用于计算机视觉和自然语言处理。

音频识别推理

音频识别 (AR) 引擎可用于识别和辨别声音模式。例如，从录音中识别鸟的类型和物种，区分音乐和歌手的声音，或检测指示建筑物破损的异常声音。为了识别感兴趣的声音，AR 引擎通过 4 个阶段处理音频

  1. 文件验证：AR 引擎验证输入音频文件。
  2. 特征提取：从音频文件中的每个片段中提取特征。
  3. 推理：LibTorch 使用 CPU 或加速器执行推理。在我们的案例中，是在 Elastic Cloud Compute (EC2) 实例上使用 Intel 处理器。
  4. 后处理：后处理模型解码结果并计算用于将推理输出转换为标签或转录的得分。

在这 4 个步骤中，推理的计算量最大，根据模型复杂性，可能占用管道处理时间的 50%。这意味着此阶段的任何优化都会对整个管道产生重大影响。

用并发性优化音频识别引擎……并非易事

此处理管道的目标是通过处理将音频片段提取为标签或转录。输入数据是一个由几个短声音片段（图 1 中的 S1 到 S6）组成的音频文件。输出数据对应于按时间戳排序的标签或转录。

Figure 1: Example audio file with segment boundaries

图 1：带有片段边界的示例音频文件

每个片段都可以独立地、无序地处理。这为并发和并行处理片段提供了机会，以优化整体推理吞吐量并最大限度地利用资源。

实例上的并行化可以通过多线程（pThreads、std::threads、OpenMP）或多进程实现。多线程相对于多进程的优势在于能够使用共享内存。它使开发人员能够通过在线程之间共享数据来最大程度地减少线程之间的数据重复；在我们的案例中，是 AR 模型（图 2）。此外，内存的减少使我们能够通过增加引擎线程的数量来并行运行更多管道，以利用我们 Amazon EC2 实例上的所有 vCPU（在我们的案例中为 c5.4xlarge，它提供 16 个 vCPU）。理论上，我们期望我们的 AR 引擎因此获得更高的硬件利用率和更高的吞吐量。

Figure 2: Multi-threaded AR Engine

图 2：多线程 AR 引擎

但我们发现这些假设是错误的。事实上，我们发现增加应用程序的线程数会导致每个音频片段的端到端延迟增加，并导致引擎吞吐量下降。例如，将并发从 1 个线程增加到 5 个线程导致延迟增加了 4 倍，这对吞吐量下降产生了成比例的影响。事实上，指标显示，在管道中，仅推理阶段的延迟就比其单线程基线高出 3 倍。

使用分析器，我们发现 CPU 空转时间增加，这可能是由于 CPU 过度订阅，从而影响系统和应用程序性能。鉴于我们对应用程序多线程实现的控制，我们选择深入研究堆栈，并识别与 LibTorch 默认设置的潜在冲突。

深入探讨 LibTorch 的多线程及其对并发性的影响

LibTorch 在 CPU 上用于推理的并行实现基于全局线程池。实现的示例是 Inter-op 和 intra-op 并行性，可以根据模型的属性进行选择。在这两种情况下，都可以设置每个线程池中的线程数以优化延迟和吞吐量。

为了测试 LibTorch 的并行默认实现设置是否对我们的推理延迟产生了反作用，我们在一个 16 vCPU 的机器上运行了一个实验，使用一个 35 分钟的音频文件，并将 LibTorch 进程间线程数保持在 1（因为我们的模型没有使用进程间线程池）。我们收集了以下数据，如图 3 和图 4 所示。

Figure 3: CPU Utilization for different number of engine threads

图 3：不同引擎线程数的 CPU 利用率

Figure 4: Processing times for different number of engine threads

图 4：不同引擎线程数的处理时间

图 4 中的执行时间是处理给定音频文件中所有片段的端到端处理时间。我们有 4 种不同的 LibTorch 进程内线程配置，分别是 1、4、8、16，并且我们针对每种进程内线程 LibTorch 配置将引擎线程数从 1 更改为 16。正如我们在图 3 中所见，对于所有 LibTorch 进程内线程配置，CPU 利用率随引擎线程数的增加而增加。但正如我们在图 4 中所见，CPU 利用率的增加并未转化为更低的执行时间。我们发现，除了一个案例之外，随着引擎线程数的激增，执行时间也随之增加。唯一的例外是进程内线程池大小为 1 的情况。

解决全局线程池问题

使用过多的线程和全局线程池会导致性能下降并引起过度订阅问题。如果不禁用LibTorch 全局线程池，就很难达到多进程引擎的性能。

禁用 LibTorch 全局线程池就像将 intra-op/inter-op 并行线程设置为 1 一样简单，如下所示

at::set_num_threads(1)           // Disables the intraop thread pool.
at::set_num_interop_threads(1). // Disables the interop thread pool.

如图 4 所示，当 LibTorch 全局线程池被禁用时，测量到的处理时间最短。

此解决方案在多种情况下提高了 AR 引擎的吞吐量。然而，在评估大型数据集（负载测试中音频文件长于 2 小时）时，我们发现引擎的内存占用量逐渐开始增加。

优化内存使用

我们使用时长两小时的音频文件对系统进行了负载测试，发现观察到的内存增加是多线程 LibTorch 推理中内存碎片化的结果。我们通过使用 jemalloc 解决了这个问题，jemalloc 是一个通用 malloc(3) 实现，它强调避免碎片化和可扩展的并发支持。使用 jemalloc 后，我们的峰值内存使用量平均下降了 34%，平均内存使用量下降了 53%。

Figure 5: Memory usage over time using the same input file with and without jemalloc

图 5：使用相同输入文件时有无 jemalloc 的内存使用随时间变化

总结

为了优化基于 LibTorch 的多线程推理引擎的性能，我们建议验证 LibTorch 中是否存在过度订阅问题。在我们的案例中，多线程引擎中的所有线程都共享 LibTorch 全局线程池，这导致了过度订阅问题。通过禁用全局线程池解决了这个问题：我们将进程间和进程内全局线程池的线程数设置为 1，从而禁用了它们。为了优化多线程引擎的内存，我们建议使用 Jemalloc 作为内存分配工具，而不是默认的 malloc 函数。