原文： https://pytorch.org/docs/stable/notes/cpu_threading_torchscript_inference.html

PyTorch 允許在 TorchScript 模型推斷期間使用多個 CPU 線程。 下圖顯示了在典型應(yīng)用程序中可以找到的不同級別的并行性：

一個或多個推理線程在給定的輸入上執(zhí)行模型的前向傳遞。 每個推理線程都調(diào)用一個 JIT 解釋器，該解釋器逐一執(zhí)行內(nèi)聯(lián)模型的操作。 模型可以利用fork TorchScript 原語來啟動異步任務(wù)。 一次分叉多個操作將導(dǎo)致并行執(zhí)行任務(wù)。 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 使用單個線程池實現(xiàn)操作間的并行性，該線程池由應(yīng)用程序過程中分叉的所有推理任務(wù)共享。

除了操作間并行性之外，PyTorch 還可以在操作內(nèi)部利用多個線程(<cite>操作內(nèi)并行性</cite>）。 在許多情況下，這可能很有用，包括大張量上的元素操作，卷積，GEMM，嵌入查找等。

構(gòu)建選項

PyTorch 使用內(nèi)部的 ATen 庫來實現(xiàn)操作。 除此之外，PyTorch 還可以通過支持 MKL 和 MKL-DNN 等外部庫來構(gòu)建，以加快 CPU 的計算速度。

ATen，MKL 和 MKL-DNN 支持操作內(nèi)并行，并依靠以下并行庫來實現(xiàn)它：

• OpenMP -廣泛用于外部庫中的標準(和庫，通常隨編譯器一起提供）；
• TBB -針對基于任務(wù)的并行性和并發(fā)環(huán)境優(yōu)化的更新并行化庫。

過去，OpenMP 已被許多庫使用。 以相對容易使用和支持基于循環(huán)的并行性和其他原語而聞名。 同時，OpenMP 與該應(yīng)用程序使用的其他線程庫之間的良好互操作性并不為人所知。 特別是，OpenMP 不保證在應(yīng)用程序中將使用單個每個進程的內(nèi)部操作線程池。 相反，兩個不同的互操作線程將可能使用不同的 OpenMP 線程池進行互操作。 這可能會導(dǎo)致應(yīng)用程序使用大量線程。

TBB 在外部庫中使用的程度較小，但同時針對并發(fā)環(huán)境進行了優(yōu)化。 PyTorch 的 TBB 后端保證了應(yīng)用程序中運行的所有操作都使用一個單獨的，按進程的單個進程內(nèi)線程池。

根據(jù)使用情況，可能會發(fā)現(xiàn)一個或另一個并行化庫在其應(yīng)用程序中是更好的選擇。

PyTorch 允許通過以下構(gòu)建選項來選擇構(gòu)建時 ATen 和其他庫使用的并行化后端：

|

圖書館

|

構(gòu)建選項

|

價值觀

|

筆記

| | --- | --- | --- | --- | | en | ATEN_THREADING | OMP(默認），TBB | | | MKL | MKL_THREADING | (相同） | 要啟用 MKL，請使用BLAS=MKL | | MKL-DNN | MKLDNN_THREADING | (same) | 要啟用 MKL-DNN，請使用USE_MKLDNN=1 |

強烈建議不要在一個內(nèi)部版本中混用 OpenMP 和 TBB。

以上任何TBB值都需要USE_TBB=1構(gòu)建設(shè)置(默認值：OFF）。 OpenMP 并行性需要單獨的設(shè)置USE_OPENMP=1(默認值：ON）。

運行時 API

以下 API 用于控制線程設(shè)置：

|

并行類型

|

設(shè)定值

|

Notes

| | --- | --- | --- | | 互操作并行 | at::set_num_interop_threads和at::get_num_interop_threads(C ++）set_num_interop_threads和get_num_interop_threads(Python， torch 模塊） | set*功能只能在啟動期間，實際操作員運行之前被調(diào)用一次；默認線程數(shù)：CPU 內(nèi)核數(shù)。 | | 幀內(nèi)并行 | at::set_num_threads，at::get_num_threads(C ++）set_num_threads，get_num_threads(Python， torch 模塊）環(huán)境變量：OMP_NUM_THREADS和MKL_NUM_THREADS |

對于操作內(nèi)并行設(shè)置，at::set_num_threads，torch.set_num_threads始終優(yōu)先于環(huán)境變量，MKL_NUM_THREADS變量優(yōu)先于OMP_NUM_THREADS。

注意

parallel_info實用程序可打印有關(guān)線程設(shè)置的信息，并可用于調(diào)試。 在 Python 中，也可以通過torch.__config__.parallel_info()調(diào)用獲得類似的輸出。