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