• 文档 >
  • CPU 线程化与 TorchScript 推理
快捷方式

CPU 线程化与 TorchScript 推理¶

PyTorch 允许在 TorchScript 模型推理期间使用多个 CPU 线程。下图展示了典型应用中可以找到的不同级别的并行性

../_images/cpu_threading_torchscript_inference.svg

一个或多个推理线程对给定输入执行模型的正向传播。每个推理线程调用一个 JIT 解释器,该解释器内联逐个执行模型的算子。模型可以使用 fork TorchScript 原语来启动异步任务。一次性分叉多个操作会产生一个并行执行的任务。fork 算子返回一个 Future 对象,该对象稍后可用于同步,例如

@torch.jit.script
def compute_z(x):
    return torch.mm(x, self.w_z)

@torch.jit.script
def forward(x):
    # launch compute_z asynchronously:
    fut = torch.jit._fork(compute_z, x)
    # execute the next operation in parallel to compute_z:
    y = torch.mm(x, self.w_y)
    # wait for the result of compute_z:
    z = torch.jit._wait(fut)
    return y + z

PyTorch 使用一个线程池来实现算子间并行性,该线程池由应用进程内所有分叉的推理任务共享。

除了算子间并行性之外,PyTorch 还可以在算子内部利用多个线程(算子内并行性)。这在许多情况下非常有用,包括大型张量的元素级算子、卷积、GEMM、嵌入查找等。

构建选项¶

PyTorch 使用内部 ATen 库来实现算子。此外,PyTorch 还可以构建以支持外部库,例如 MKLMKL-DNN,以加快 CPU 上的计算速度。

ATen、MKL 和 MKL-DNN 支持算子内并行性,并依赖以下并行库来实现它

  • OpenMP - 一个标准(和库,通常随编译器一起提供),在外部库中广泛使用;

  • TBB - 一个较新的并行化库,针对基于任务的并行性和并发环境进行了优化。

OpenMP 历来被大量库使用。它以相对易用和支持基于循环的并行性及其他原语而闻名。

TBB 在外部库中使用的程度较低,但同时针对并发环境进行了优化。PyTorch 的 TBB 后端保证应用程序中所有运行的算子都使用一个单独的、单一的、每进程的算子内线程池。

根据用例的不同,可能会发现某个并行化库在其应用中是更好的选择。

PyTorch 允许在构建时通过以下构建选项选择 ATen 和其他库使用的并行化后端

构建选项

备注

ATen

ATEN_THREADING

OMP (默认), TBB

MKL

MKL_THREADING

(同上)

要启用 MKL,请使用 BLAS=MKL

MKL-DNN

MKLDNN_CPU_RUNTIME

(同上)

要启用 MKL-DNN,请使用 USE_MKLDNN=1

建议不要在同一构建中混合使用 OpenMP 和 TBB。

上述任何 TBB 值都需要 USE_TBB=1 构建设置(默认:OFF)。对于 OpenMP 并行性,需要单独的设置 USE_OPENMP=1(默认:ON)。

运行时 API¶

以下 API 用于控制线程设置

并行类型

设置

备注

算子间并行性

at::set_num_interop_threads, at::get_num_interop_threads (C++)

set_num_interop_threads, get_num_interop_threads (Python, torch 模块)

默认线程数:CPU 核心数。

算子内并行性

at::set_num_threads, at::get_num_threads (C++) set_num_threads, get_num_threads (Python, torch 模块)

环境变量:OMP_NUM_THREADSMKL_NUM_THREADS

对于算子内并行性设置,at::set_num_threadstorch.set_num_threads 始终优先于环境变量,MKL_NUM_THREADS 变量优先于 OMP_NUM_THREADS

调整线程数¶

以下简单脚本展示了矩阵乘法的运行时如何随线程数变化

import timeit
runtimes = []
threads = [1] + [t for t in range(2, 49, 2)]
for t in threads:
    torch.set_num_threads(t)
    r = timeit.timeit(setup = "import torch; x = torch.randn(1024, 1024); y = torch.randn(1024, 1024)", stmt="torch.mm(x, y)", number=100)
    runtimes.append(r)
# ... plotting (threads, runtimes) ...

在具有 24 个物理 CPU 核心(Xeon E5-2680,基于 MKL 和 OpenMP 构建)的系统上运行该脚本,结果如下所示的运行时

../_images/cpu_threading_runtimes.svg

调整算子内和算子间线程数时应考虑以下事项

  • 选择线程数时,需要避免过度订阅(使用过多线程会导致性能下降)。例如,在使用大型应用线程池或高度依赖算子间并行性的应用中,禁用算子内并行性(即通过调用 set_num_threads(1))可能是一个选项;

  • 在典型应用中,可能会遇到延迟(处理推理请求所花费的时间)和吞吐量(单位时间内完成的工作量)之间的权衡。调整线程数是调整这种权衡的有用工具。例如,在延迟关键型应用中,可能希望增加算子内线程数以尽快处理每个请求。同时,算子的并行实现可能会增加额外的开销,从而增加单个请求完成的工作量,进而降低总体吞吐量。

警告

OpenMP 不保证应用程序中会使用单个每进程算子内线程池。相反,两个不同的应用程序线程或算子间线程可能使用不同的 OpenMP 线程池来执行算子内工作。这可能导致应用程序使用大量线程。在 OpenMP 情况下,多线程应用程序中需要格外注意调整线程数以避免过度订阅。

注意

预构建的 PyTorch 版本是使用 OpenMP 支持编译的。

注意

parallel_info 工具打印有关线程设置的信息,可用于调试。在 Python 中,通过调用 torch.__config__.parallel_info() 也可以获得类似的输出。

文档

访问 PyTorch 全面的开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源