作者：Himalay Mohanlal Joriwal, Pierre-Yves Aquilanti, Vivek Govindan, Hamid Shojanazeri, Ankith Gunapal, Tristan Rice

概要

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

音频识别推理

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

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

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

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

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

Figure 1: Example audio file with segment boundaries

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

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

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

Figure 2: Multi-threaded AR Engine

图 2：多线程 AR 引擎

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

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

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

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

为了测试 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 全局线程池只需将操作内/跨操作并行线程数设置为 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 解决了这个问题，它是一个通用的 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 函数。