CPU 執行緒與 TorchScript 推論

PyTorch 允許在 TorchScript 模型推論期間使用多個 CPU 執行緒。下圖顯示了一個典型應用程式中可能存在的不同層次的平行處理

一個或多個推論執行緒對給定的輸入執行模型的正向傳遞。每個推論執行緒會呼叫 JIT 直譯器，該直譯器會逐一執行模型的操作。模型可以使用 fork TorchScript 原始指令來啟動非同步任務。一次 Fork 多個操作會產生一個並行執行的任務。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 使用單一執行緒池進行操作間平行處理，此執行緒池由應用程式進程中 fork 的所有推論任務共享。

除了操作間平行處理之外，PyTorch 還可以在操作內利用多個執行緒（操作內平行處理）。這在許多情況下都很有用，包括大型張量上的元素級操作、卷積、GEMM、嵌入式查找等。

建置選項

PyTorch 使用內部的 ATen 函式庫來實現操作。除此之外，PyTorch 也可以建置為支援外部函式庫，例如 MKL 和 MKL-DNN，以加速 CPU 上的計算。

ATen、MKL 和 MKL-DNN 支援操作內平行處理，並依賴以下平行處理函式庫來實現它

• OpenMP - 一個標準 (也是一個函式庫，通常隨編譯器一起提供)，廣泛應用於外部函式庫；

• TBB - 一個較新的平行化函式庫，針對基於任務的平行處理和並行環境進行了優化。

OpenMP 歷史上已被大量函式庫使用。 它以相對易於使用、支持基於迴圈的平行處理和其他基本元件而聞名。

TBB 在外部函式庫中的使用程度較低，但同時針對並行環境進行了優化。 PyTorch 的 TBB 後端保證應用程式中運行的所有運算 (op) 都使用獨立、單一的、每個進程的運算內 (intra-op) 執行緒池。

根據使用案例，您可能會發現一個或另一個平行化函式庫更適合您的應用程式。

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 用於控制執行緒設定

平行處理類型	設定	備註
運算間 (Inter-op) 平行處理	at::set_num_interop_threads, at::get_num_interop_threads (C++) set_num_interop_threads, get_num_interop_threads (Python, torch 模組)	預設執行緒數：CPU 核心數。
運算內 (Intra-op) 平行處理	at::set_num_threads, at::get_num_threads (C++) set_num_threads, get_num_threads (Python, torch 模組) 環境變數: OMP_NUM_THREADS 和 MKL_NUM_THREADS

對於運算內平行處理設定，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 工具會印出有關執行緒設定的資訊，可用於除錯。 也可以通過在 Python 中調用 torch.__config__.parallel_info() 來獲得類似的輸出。