pytorch編程之CPU 線程和 TorchScript 推斷
PyTorch 允許在 TorchScript 模型推斷期間使用多個(gè) CPU 線程。下圖顯示了在典型應(yīng)用程序中可以找到的不同級(jí)別的并行性：
一個(gè)或多個(gè)推理線程在給定的輸入上執(zhí)行模型的前向傳遞。每個(gè)推理線程都調(diào)用一個(gè) JIT 解釋器，該解釋器逐一執(zhí)行內(nèi)聯(lián)模型的操作。模型可以利用fork?TorchScript 原語(yǔ)來(lái)啟動(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 + zCopyPyTorch 使用單個(gè)線程池實(shí)現(xiàn)操作間的并行性，該線程池由應(yīng)用程序過(guò)程中分叉的所有推理任務(wù)共享。
除了操作間并行性之外，PyTorch 還可以在操作內(nèi)部利用多個(gè)線程(操作內(nèi)并行性）。在許多情況下，這可能很有用，包括大張量上的元素操作，卷積，GEMM，嵌入查找等。
構(gòu)建選項(xiàng)
PyTorch 使用內(nèi)部的 ATen 庫(kù)來(lái)實(shí)現(xiàn)操作。除此之外，PyTorch 還可以通過(guò)支持 MKL 和 MKL-DNN 等外部庫(kù)來(lái)構(gòu)建，以加快 CPU 的計(jì)算速度。
ATen，MKL 和 MKL-DNN 支持操作內(nèi)并行，并依靠以下并行庫(kù)來(lái)實(shí)現(xiàn)它：
OpenMP -廣泛用于外部庫(kù)中的標(biāo)準(zhǔn)(和庫(kù)，通常隨編譯器一起提供）；
TBB -針對(duì)基于任務(wù)的并行性和并發(fā)環(huán)境優(yōu)化的更新并行化庫(kù)。
過(guò)去，OpenMP 已被許多庫(kù)使用。以相對(duì)容易使用和支持基于循環(huán)的并行性和其他原語(yǔ)而聞名。同時(shí)，OpenMP 與該應(yīng)用程序使用的其他線程庫(kù)之間的良好互操作性并不為人所知。特別是，OpenMP 不保證在應(yīng)用程序中將使用單個(gè)每個(gè)進(jìn)程的內(nèi)部操作線程池。相反，兩個(gè)不同的互操作線程將可能使用不同的 OpenMP 線程池進(jìn)行互操作。這可能會(huì)導(dǎo)致應(yīng)用程序使用大量線程。
TBB 在外部庫(kù)中使用的程度較小，但同時(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è)并行化庫(kù)在其應(yīng)用程序中是更好的選擇。
PyTorch 允許通過(guò)以下構(gòu)建選項(xiàng)來(lái)選擇構(gòu)建時(shí) ATen 和其他庫(kù)使用的并行化后端：
|
圖書(shū)館
|
構(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變量?jī)?yōu)先于OMP_NUM_THREADS。
注意
parallel_info實(shí)用程序可打印有關(guān)線程設(shè)置的信息，并可用于調(diào)試。在 Python 中，也可以通過(guò)torch.__config__.parallel_info()調(diào)用獲得類似的輸出。