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):
# 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还可以在操作内利用多个线程 (intra-op parallelism)。这在许多情况下都很有用，包括对大型张量进行逐元素操作、卷积、GEMM、嵌入查找等。
构建选项¶
PyTorch 使用内部的 ATen 库来实现操作。除此之外，PyTorch 还可以构建对外部库的支持，如 MKL 和 MKL-DNN，以加速 CPU 上的计算。
ATen、MKL 和 MKL-DNN 支持操作内并行，并依赖以下并行库来实现它：
OpenMP 历史上已被大量库使用。它以相对易于使用和支持基于循环的并行性和其他原语而闻名。
TBB 在外部库中的使用较少，但同时针对并发环境进行了优化。PyTorch 的 TBB 后端保证了有一个独立的、每个进程内的单个操作线程池，用于应用程序中运行的所有操作。
根据使用情况，人们可能会在他们的应用程序中发现一个或另一个并行化库是更好的选择。
PyTorch 允许在构建时通过以下构建选项选择 ATen 和其他库使用的并行化后端：
库 |
构建选项 |
值 |
注释 |
|---|---|---|---|
ATen |
|
|
|
MKL |
|
(相同) |
要启用MKL，请使用 |
MKL-DNN |
|
(相同) |
要启用MKL-DNN，请使用 |
建议不要在一个构建中混合使用OpenMP和TBB。
上述任何 TBB 个值都需要 USE_TBB=1 个构建设置（默认：OFF）。
OpenMP 并行性需要单独的设置 USE_OPENMP=1（默认：ON）。
运行时API¶
以下API用于控制线程设置：
并行类型 |
设置 |
注释 |
|---|---|---|
进程间并行性 |
|
默认线程数：CPU核心数。 |
操作内并行性 |
环境变量:
|
对于操作内并行设置，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的构建）的系统上运行脚本，得到以下运行时间：
在调整内部和跨操作线程数量时，应考虑以下因素：
在选择线程数量时，需要避免oversubscription（使用过多线程会导致性能下降）。例如，在一个使用大型应用程序线程池或高度依赖于跨操作并行的应用程序中，可能会发现禁用操作内并行性是一个可能的选项（即通过调用
set_num_threads(1)）；在典型的应用中，人们可能会遇到在latency（处理推理请求所花费的时间）和throughput（单位时间内完成的工作量）之间进行权衡。调整线程数量可以成为以某种方式调整这种权衡的一个有用工具。例如，在延迟关键的应用程序中，人们可能希望增加每个操作的线程数，以便尽可能快速地处理每个请求。同时，并行实现的操作可能会增加额外的开销，从而增加每个单独请求的工作量，因此降低整体吞吐量。
警告
OpenMP 并不能保证应用程序中会使用单个每进程的内部操作线程池。相反，两个不同的应用程序或跨操作线程可能会使用不同的 OpenMP 线程池来执行内部操作任务。这可能导致应用程序使用大量线程。在多线程应用程序中，特别是在使用 OpenMP 的情况下，需要特别注意调整线程数量以避免过度订阅。
注意
预编译的PyTorch版本是带有OpenMP支持进行编译的。
注意
parallel_info 实用程序打印关于线程设置的信息，可用于调试。
类似输出也可以通过Python中的 torch.__config__.parallel_info() 调用来获得。