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	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++) set_num_interop_threads、get_num_interop_threads (Python，torch 模組)	預設線程數：CPU 核心數。
操作內平行	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() 獲得類似的輸出。