CPU 线程和 TorchScript 推理

PyTorch 允许在 TorchScript 模型推理期间使用多个 CPU 线程。 下图显示了在 典型应用：

一个或多个推理线程对给定的输入执行模型的正向传递。 每个推理线程都会调用一个 JIT 解释器来执行运算 的模型内联。模型可以使用 TorchScript primitive 启动异步任务。一次分叉多个操作 导致并行执行的任务。该运算符返回一个对象，该对象可用于稍后同步，例如：forkforkFuture

@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 还可以利用多个线程 在 OPS 中（OP 内并行）。这在许多情况下都很有用， 包括对大型张量、卷积、GEMM、嵌入的元素级运算 查找和其他。

构建选项

PyTorch 使用内部 ATen 库来实现运算。除此之外， PyTorch 也可以在支持外部库（如 MKL 和 MKL-DNN）的情况下构建。 以加快 CPU 上的计算速度。

ATen、MKL 和 MKL-DNN 支持操作内并行性，并依赖于 遵循并行化库来实现它：

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

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

OpenMP 历史上已被大量库使用。众所周知 相对易于使用并支持基于循环的并行性和其他基元。

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

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

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。

上述任何值都需要 build 设置（默认值：OFF）。 OpenMP 并行性需要单独的设置（默认值：ON）。TBBUSE_TBB=1USE_OPENMP=1

运行时 API

以下 API 用于控制线程设置：

平行度类型	设置	笔记
互操作性	at::set_num_interop_threads、（C++）at::get_num_interop_threads set_num_interop_threads， （Python， module）get_num_interop_threadstorch	Default number of threads：CPU 核心数。
操作内并行性	at::set_num_threads， （C++） ， （Python， 模块）at::get_num_threadsset_num_threadsget_num_threadstorch 环境变量： 和OMP_NUM_THREADSMKL_NUM_THREADS

对于操作内并行度设置，始终优先 在环境变量中， variable 优先于 。at::set_num_threadstorch.set_num_threadsMKL_NUM_THREADSOMP_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 的构建）的系统上运行脚本会产生以下运行时：

在调整操作内线程和操作间线程的数量时，应考虑以下注意事项：

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

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

警告

OpenMP 不保证单个每进程内部操作线程 pool 将在应用程序中使用。相反，两个不同的应用程序或互操作 线程可以使用不同的 OpenMP 线程池进行内部操作工作。 这可能会导致应用程序使用大量线程。 在调整线程数时需要格外小心，以避免 OpenMP 案例中多线程应用程序中的超额订阅。

注意

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

注意

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