CPU 线程和 TorchScript 推理

PyTorch 允许在 TorchScript 模型推理期间使用多个 CPU 线程。下图显示了一个典型应用程序中可能存在的不同级别的并行性

一个或多个推理线程在给定的输入上执行模型的前向传递。每个推理线程调用一个 JIT 解释器，该解释器逐个内联执行模型的操作。模型可以利用 fork TorchScript 原语来启动异步任务。一次 fork 多个操作会导致任务并行执行。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 为 inter-op 并行性使用单个线程池，此线程池由应用程序进程中 fork 的所有推理任务共享。

除了 inter-op 并行性之外，PyTorch 还可以在操作内部利用多个线程（intra-op 并行性）。这在许多情况下都很有用，包括大型张量上的逐元素操作、卷积、GEMM、嵌入查找等。

构建选项

PyTorch 使用内部 ATen 库来实现操作。除此之外，PyTorch 还可以构建为支持外部库，例如 MKL 和 MKL-DNN，以加速 CPU 上的计算。

ATen、MKL 和 MKL-DNN 支持 intra-op 并行性，并依赖以下并行库来实现它

• OpenMP - 一种标准（以及一个库，通常随编译器一起提供），广泛用于外部库；

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

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

TBB 在外部库中的使用程度较低，但同时针对并发环境进行了优化。PyTorch 的 TBB 后端保证应用程序中运行的所有操作都使用单独的、单进程的 intra-op 线程池。

根据用例，人们可能会发现一个或另一个并行库在他们的应用程序中是更好的选择。

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 用于控制线程设置

并行类型	设置	注释
Inter-op 并行性	at::set_num_interop_threads, at::get_num_interop_threads (C++) set_num_interop_threads, get_num_interop_threads (Python, torch 模块)	默认线程数：CPU 核心数。
Intra-op 并行性	at::set_num_threads, at::get_num_threads (C++) set_num_threads, get_num_threads (Python, torch 模块) 环境变量：OMP_NUM_THREADS 和 MKL_NUM_THREADS

对于 intra-op 并行性设置，at::set_num_threads, torch.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 构建）的系统上运行该脚本会得到以下运行时

在调整 intra-op 和 inter-op 线程数时，应考虑以下事项

• 在选择线程数时，需要避免过度订阅（使用过多线程，导致性能下降）。例如，在大量使用应用程序线程池或严重依赖 inter-op 并行性的应用程序中，人们可能会发现禁用 intra-op 并行性是一种可能的选择（即，通过调用 set_num_threads(1)）；

• 在典型的应用程序中，人们可能会遇到延迟（处理推理请求所花费的时间）和吞吐量（单位时间内完成的工作量）之间的权衡。调整线程数可以成为以某种方式调整这种权衡的有用工具。例如，在对延迟敏感的应用程序中，人们可能希望增加 intra-op 线程数，以尽可能快地处理每个请求。同时，操作的并行实现可能会增加额外的开销，从而增加每个请求完成的工作量，从而降低整体吞吐量。

警告

OpenMP 不保证应用程序中将使用单个进程内 intra-op 线程池。相反，两个不同的应用程序或 inter-op 线程可能对 intra-op 工作使用不同的 OpenMP 线程池。这可能会导致应用程序使用大量线程。在 OpenMP 情况下，在多线程应用程序中需要格外小心地调整线程数以避免过度订阅。

注意

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

注意

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