CPU 线程和 TorchScript 推理

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

一个或多个推理线程在给定的输入上执行模型的前向传播。 每个推理线程调用一个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):
    # 异步启动 compute_z：
    fut = torch.jit._fork(compute_z, x)
    # 在 compute_z 的同时执行下一个操作：
    y = torch.mm(x, self.w_y)
    # 等待 compute_z 的结果：
    z = torch.jit._wait(fut)
    return y + z

PyTorch 使用一个线程池用于内部并行，这个线程池由应用程序进程中所有派生的推理任务共享。

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

构建选项

PyTorch 使用内部 ATen 库来实现操作。除此之外，PyTorch 还可以构建为支持外部库，例如 MKL 和 MKL-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构建设置（默认：关闭）。 对于OpenMP并行性，需要单独的设置USE_OPENMP=1（默认：开启）。

运行时 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_THREADS 和 MKL_NUM_THREADS

对于操作内并行设置，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)
# ... 绘图 (threads, runtimes) ...

在具有24个物理CPU核心（Xeon E5-2680，基于MKL和OpenMP构建）的系统上运行脚本，结果如下运行时间：

在调整intra-和inter-op线程数量时应考虑以下因素：

• 在选择线程数量时，需要避免过度订阅（使用过多线程，导致性能下降）。例如，在一个使用大型应用程序线程池或严重依赖互操作并行性的应用程序中，可能会发现禁用内部操作并行性作为一个可能的选项（即通过调用set_num_threads(1)）；

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

警告

OpenMP 不保证应用程序中会使用一个单一的每个进程的内部操作线程池。相反，两个不同的应用程序或交互操作线程可能会使用不同的 OpenMP 线程池来进行内部操作工作。这可能会导致应用程序使用大量的线程。在 OpenMP 情况下，需要特别注意调整线程数量，以避免多线程应用程序中的过度订阅。

注意

预构建的 PyTorch 版本已编译为支持 OpenMP。

注意

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