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

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

一個(gè)或多個(gè)推理線程在給定的輸入上執(zhí)行模型的前向傳遞。 每個(gè)推理線程都調(diào)用一個(gè) JIT 解釋器，該解釋器逐一執(zhí)行內(nèi)聯(lián)模型的操作。 模型可以利用fork TorchScript 原語來啟動(dòng)異步任務(wù)。 一次分叉多個(gè)操作將導(dǎo)致并行執(zhí)行任務(wù)。 fork運(yùn)算符返回一個(gè)future對(duì)象，該對(duì)象可用于以后進(jìn)行同步，例如：

@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 使用單個(gè)線程池實(shí)現(xiàn)操作間的并行性，該線程池由應(yīng)用程序過程中分叉的所有推理任務(wù)共享。

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

構(gòu)建選項(xiàng)

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

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

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

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

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

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

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

|

圖書館

|

構(gòu)建選項(xiàng)

|

價(jià)值觀

|

筆記

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

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

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

運(yùn)行時(shí) 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*功能只能在啟動(dòng)期間，實(shí)際操作員運(yùn)行之前被調(diào)用一次；默認(rèn)線程數(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 |

對(duì)于操作內(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實(shí)用程序可打印有關(guān)線程設(shè)置的信息，并可用于調(diào)試。 在 Python 中，也可以通過torch.__config__.parallel_info()調(diào)用獲得類似的輸出。