評價此頁
★ ★ ★ ★ ★

多程序最佳實踐#

建立日期: 2017年1月16日 | 最後更新日期: 2025年6月18日

torch.multiprocessing 是 Python 內建 multiprocessing 模組的“即插即用”替代品。它支援完全相同的操作，但對其進行了擴充套件，因此透過 multiprocessing.Queue 傳輸的所有張量，其資料將被移至共享記憶體，並且只會將一個控制代碼傳送到另一個程序。

注意

當 Tensor 被髮送到另一個程序時，Tensor 資料是共享的。如果 torch.Tensor.grad 不是 None，它也會被共享。在將一個沒有 torch.Tensor.grad 欄位的 Tensor 傳送到另一個程序後，它會建立一個標準的、特定於程序的 .grad Tensor，它不會像 Tensor 資料那樣被自動跨程序共享。

這使得實現各種訓練方法成為可能，例如 Hogwild、A3C 或任何需要非同步操作的方法。

多程序中的“毒性 fork”#

在使用帶 加速器 的多程序時，可能會出現一個稱為“毒性 fork”的已知問題。當加速器的執行時不安全且在程序 fork 之前初始化時，就會發生這種情況，從而導致子程序中的執行時錯誤。

為防止此類錯誤

  • 在 fork 子程序之前，避免在主程序中初始化加速器。

  • 使用替代的程序啟動方法，例如 spawn 或 forkserver，這確保了每個程序都能進行乾淨的初始化。

多程序中的 CUDA#

在使用 fork 啟動方法時，CUDA 執行時存在 多程序中的“毒性 fork” 中描述的限制；必須使用 spawn 或 forkserver 啟動方法才能在子程序中使用 CUDA。

注意

啟動方法可以透過建立具有 multiprocessing.get_context(...) 的上下文或直接使用 multiprocessing.set_start_method(...) 來設定。

與 CPU 張量不同，傳送程序在接收程序保留張量副本時，必須保持原始張量。這在底層實現，但要求使用者遵循最佳實踐才能使程式正確執行。例如，傳送程序必須在消費者程序仍然持有張量引用時保持活動狀態，如果消費者程序由於致命訊號異常退出，則引用計數無法挽救您。請參閱 本節。

另請參閱: 使用 nn.parallel.DistributedDataParallel 而不是 multiprocessing 或 nn.DataParallel

最佳實踐和技巧#

避免和處理死鎖#

在生成新程序時，很多事情都可能出錯，死鎖最常見的原因是後臺執行緒。如果任何執行緒持有鎖或匯入了模組，並且呼叫了 fork，那麼子程序很可能處於損壞狀態，並會導致死鎖或其他故障。請注意，即使您沒有，Python 的內建庫也會如此——無需查詢比 multiprocessing 更遠。 multiprocessing.Queue 實際上是一個非常複雜的類，它會生成多個執行緒用於序列化、傳送和接收物件，它們也可能導致上述問題。如果您發現自己處於這種情況，請嘗試使用 SimpleQueue，它不使用任何額外的執行緒。

我們正在盡最大努力讓您輕鬆使用，並確保這些死鎖不會發生，但有些事情是我們無法控制的。如果您遇到任何一段時間內無法解決的問題，請嘗試在論壇上尋求幫助，我們將看看這是否是我們能夠解決的問題。

重用透過 Queue 傳遞的緩衝區#

請記住，每次將 Tensor 放入 multiprocessing.Queue 時，它都必須被移動到共享記憶體中。如果它已經是共享的，則為無操作，否則將產生額外的記憶體複製，這會減慢整個過程。即使您有一個程序池將資料傳送到一個程序，也要讓它傳送回緩衝區——這幾乎是免費的，並且可以讓您在傳送下一批資料時避免複製。

非同步多程序訓練（例如 Hogwild）#

使用 torch.multiprocessing，可以非同步訓練模型，引數要麼一直共享，要麼定期同步。在第一種情況下，我們建議傳輸整個模型物件，而在後一種情況下，我們建議只傳輸 state_dict()。

我們建議使用 multiprocessing.Queue 在程序之間傳遞所有型別的 PyTorch 物件。例如，在使用 fork 啟動方法時，可以繼承已經在共享記憶體中的張量和儲存，但這非常容易出錯，應謹慎使用，並且僅供高階使用者使用。Queue，即使它們有時解決方案不那麼優雅，但在所有情況下都可以正常工作。

警告

您應該注意不要有全域性語句，這些語句沒有用 if __name__ == '__main__' 進行保護。如果使用了不同於 fork 的啟動方法，它們將在所有子程序中執行。

Hogwild#

Hogwild 的具體實現可以在 examples 倉庫 中找到，但為了展示程式碼的整體結構，下面還有一個最小示例。

import torch.multiprocessing as mp
from model import MyModel

def train(model):
    # Construct data_loader, optimizer, etc.
    for data, labels in data_loader:
        optimizer.zero_grad()
        loss_fn(model(data), labels).backward()
        optimizer.step()  # This will update the shared parameters

if __name__ == '__main__':
    num_processes = 4
    model = MyModel()
    # NOTE: this is required for the ``fork`` method to work
    model.share_memory()
    processes = []
    for rank in range(num_processes):
        p = mp.Process(target=train, args=(model,))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

多程序中的 CPU#

不恰當的多程序可能導致 CPU 過載，使不同的程序爭奪 CPU 資源，從而導致效率低下。

本教程將解釋什麼是 CPU 過載以及如何避免它。

CPU 過載#

CPU 過載是一個技術術語，指的是分配給系統的 vCPU 總數超過硬體可用 vCPU 總數的情況。

這會導致 CPU 資源嚴重爭用。在這種情況下，程序之間會頻繁切換，這會增加程序切換的開銷並降低整體系統效率。

在 示例倉庫 中找到的 Hogwild 實現的程式碼示例中，可以看到 CPU 過載。

在 CPU 上使用 4 個程序執行訓練示例時，命令如下：

python main.py --num-processes 4

假設機器上有 N 個可用的 vCPU，執行上述命令將生成 4 個子程序。每個子程序將為自己分配 N 個 vCPU，導致需要 4*N 個 vCPU。但是，機器上只有 N 個可用的 vCPU。因此，不同的程序將爭奪資源，導致頻繁的程序切換。

以下觀察結果表明存在 CPU 過載：

  1. 高 CPU 利用率：使用 htop 命令，您可以觀察到 CPU 利用率持續很高，通常達到或超過其最大容量。這表明對 CPU 資源的需求超過了可用的物理核心，導致程序之間爭奪 CPU 時間。

  2. 頻繁的上下文切換導致系統效率低下：在 CPU 過載的情況下，程序會爭奪 CPU 時間，作業系統需要快速在不同程序之間切換以公平地分配資源。這種頻繁的上下文切換會增加開銷並降低整體系統效率。

避免 CPU 過載#

避免 CPU 過載的一個好方法是進行適當的資源分配。確保併發執行的程序或執行緒數量不超過可用的 CPU 資源。

在這種情況下，一種解決方案是指定子程序中的適當執行緒數。這可以透過在子程序中使用 torch.set_num_threads(int) 函式來設定每個程序的執行緒數來實現。

假設機器上有 N 個 vCPU，並且將生成 M 個程序，那麼每個程序使用的最大 num_threads 值將是 floor(N/M)。為了避免 mnist_hogwild 示例中的 CPU 過載，需要對 示例倉庫 中的 train.py 檔案進行以下更改。

def train(rank, args, model, device, dataset, dataloader_kwargs):
    torch.manual_seed(args.seed + rank)

    #### define the num threads used in current sub-processes
    torch.set_num_threads(floor(N/M))

    train_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs)

    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
    for epoch in range(1, args.epochs + 1):
        train_epoch(epoch, args, model, device, train_loader, optimizer)

使用 torch.set_num_threads(floor(N/M)) 為每個程序設定 num_thread。其中 N 是可用 vCPU 的數量，M 是選擇的程序數。合適的 num_thread 值將因具體任務而異。但是，作為一般指南，num_thread 的最大值應為 floor(N/M) 以避免 CPU 過載。在 mnist_hogwild 訓練示例中，在避免 CPU 過載後，可以實現 30 倍的效能提升。

在此頁面上
在 GitHub 上編輯
顯示原始碼
PyTorch 庫