# CPU 多线程和 TorchScript 推理
PyTorch 允许在 TorchScript 模型推理过程中使用多个 CPU 线程。下图展示了典型应用中不同级别的并行性：
[]({BASE_RAW_UPLOAD_URL}/pytorch-doc-2.5/8df78fa0159321538b2e2a438f6cae52.svg)
一个或多个推理线程在给定输入上执行模型的前向传播。每个推理线程调用一个 JIT 解释器，依次执行模型的操作。模型可以使用 `fork` TorchScript 原语来发起异步任务。同时启动多个操作会导致任务并行执行。`fork` 操作符返回一个 `Future` 对象，可用于稍后的同步，例如：
```plain
@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 使用单个线程池来实现操作间的并行性（inter-op parallelism），该线程池由应用程序进程内的所有推理任务共享。
除了操作间的并行性外，PyTorch 还可以在操作内部利用多个线程（intra-op parallelism）。这在许多情况下非常有用，例如对大张量进行逐元素操作、卷积、GEMM、嵌入查找等。
## 构建选项
PyTorch 使用内部的 ATen 库来实现操作。除此之外，PyTorch 还可以构建对外部库的支持，例如 [MKL](https://software.intel.com/en-us/mkl) 和 [MKL-DNN](https://github.com/intel/mkl-dnn)，以加速 CPU 上的计算。
ATen、MKL 和 MKL-DNN 支持操作内的并行性，并依赖以下并行化库来实现这些功能：
* [OpenMP](https://www.openmp.org/) - 一个标准（也是一个库，通常随编译器一起提供），在外部库中广泛使用；
* [TBB](https://github.com/intel/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` 构建设置（默认：OFF）。对于 OpenMP 并行性，需要单独的设置 `USE_OPENMP=1`（默认：ON）。
## 运行时 API
以下 API 用于控制线程设置：
| 并行类型 | 设置 | 备注 |
| --- | --- | --- |
| 操作间并行 | `at::set_num_interop_threads`, `at::get_num_interop_threads` (C++)<br><br>`set_num_interop_threads`, `get_num_interop_threads` (Python, [`torch`](../torch.html#module-torch) 模块) | 默认线程数：CPU 核心数。 |
| 操作内并行 | `at::set_num_threads`, `at::get_num_threads` (C++) `set_num_threads`, `get_num_threads` (Python, [`torch`](../torch.html#module-torch) 模块)<br><br>环境变量：`OMP_NUM_THREADS` 和 `MKL_NUM_THREADS` |
对于操作内并行设置，`at::set_num_threads` 和 `torch.set_num_threads` 始终优先于环境变量，`MKL_NUM_THREADS` 变量优先于 `OMP_NUM_THREADS`。
## 调整线程数
以下简单的脚本展示了矩阵乘法的运行时间如何随着线程数的变化而变化：
plain import timeit runtimes = [] threads = [1] + [t for t in range(2, 49, 2)] for t in threads: torch.setnumthreads(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 的构建）的系统上运行该脚本，得到的运行时间如下：
在调整操作内和操作间线程数时，应考虑以下几点：
-
选择线程数时，需要避免过度分配（使用过多的线程会导致性能下降）。例如，在使用大量应用线程池或严重依赖于操作间并行性的应用程序中，可以考虑禁用操作内并行性（即调用
set_num_threads(1)）； -
在典型的应用程序中，可能会遇到延迟（处理请求的时间）和吞吐量（每单位时间完成的工作量）之间的权衡。调整线程数可以作为一种有用的工具来调整这种权衡。例如，在对延迟敏感的应用程序中，可能希望增加操作内线程数以尽可能快地处理每个请求。同时，操作的并行实现可能会增加额外的开销，从而增加单个请求的工作量，从而降低整体吞吐量。
警告
OpenMP 并不保证在应用程序中会使用单个每进程的 intra-op 线程池。相反，两个不同的应用程序或 inter-op 线程可能会使用不同的 OpenMP 线程池来处理 intra-op 任务。这可能导致应用程序使用大量线程。在多线程应用程序中，需要特别小心地调整线程数量，以避免资源过度占用。
注意
预编译的 PyTorch 版本支持 OpenMP。
注意
parallel_info 工具可以打印线程设置信息，用于调试。类似的信息也可以通过在 Python 中调用 torch.__config__.parallel_info() 获取。