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

大纲

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

音频识别推理

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

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

在这 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 引擎因此实现更高的硬件利用率和更高的吞吐量。

图 2：多线程 AR 引擎

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

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

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

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

为了测试 LibTorch 的并行默认实现设置是否对我们的推理延迟产生了反作用，我们在一个 16 vCPU 机器上运行了一个实验，使用一个 35 分钟的音频文件，保持 LibTorch 内部线程常数为 1（因为我们的模型没有使用 inter-op 线程池）。我们收集了以下数据，如图 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 全局线程池非常简单，只需将 inter-op/intra-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 全局线程池，这导致了过度订阅问题。通过禁用全局线程池解决了这个问题：我们将 interop 和 intraop 全局线程池的线程数设置为 1 来禁用它们。为了优化多线程引擎的内存，我们建议使用 Jemalloc 作为内存分配工具，而不是默认的 malloc 函数。