多程序最佳實踐¶
torch.multiprocessing 是 Python 標準庫 multiprocessing 模組的直接替代品。它支援完全相同的操作,並對其進行了擴充套件,使得所有透過 multiprocessing.Queue 傳送的 tensor 的資料都會被移動到共享記憶體中,並僅將一個控制代碼傳送給另一個程序。
注意
當一個 Tensor 被髮送到另一個程序時,該 Tensor 的資料是共享的。如果 torch.Tensor.grad 不是 None,它也會被共享。一個沒有 torch.Tensor.grad 欄位的 Tensor 被髮送到其他程序後,會在接收程序中建立一個標準的、程序專屬的 .grad Tensor,與 Tensor 資料已共享的方式不同,這個 .grad Tensor 不會自動在所有程序中共享。
這使得實現各種訓練方法成為可能,例如 Hogwild、A3C 或任何其他需要非同步操作的方法。
多程序中的 CUDA¶
CUDA 執行時不支援 fork 啟動方法;在子程序中使用 CUDA 需要使用 spawn 或 forkserver 啟動方法。
注意
啟動方法可以透過建立上下文 (`multiprocessing.get_context(...)`) 或直接使用 multiprocessing.set_start_method(...) 來設定。
與 CPU tensor 不同,只要接收程序保留了該 tensor 的副本,傳送程序就需要保留原始 tensor。這一點在底層已經實現,但要求使用者遵循最佳實踐以確保程式正確執行。例如,只要消費者程序持有 tensor 的引用,傳送程序就必須保持存活;如果消費者程序因致命訊號異常退出,引用計數機制無法拯救你。參見 本節。
另請參閱:使用 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 啟動方法時,可以繼承共享記憶體中已有的 tensor 和 storage,但這非常容易出錯,應謹慎使用,且僅限高階使用者。Queue,儘管有時不是那麼優雅的解決方案,但在所有情況下都能正常工作。
警告
你應該小心那些沒有使用 if __name__ == '__main__' 保護的全域性語句。如果使用了 fork 以外的啟動方法,它們會在所有子程序中執行。
Hogwild¶
可以在 examples 倉庫中找到一個具體的 Hogwild 實現,但為了展示程式碼的總體結構,下面也有一個最小示例
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 過度佔用(oversubscription),使得不同程序競爭 CPU 資源,從而導致效率低下。
本教程將解釋 CPU 過度佔用是什麼以及如何避免它。
CPU 過度佔用¶
CPU 過度佔用是一個技術術語,指系統中分配的 vCPU 總數超過硬體上可用的 vCPU 總數的情況。
這會導致對 CPU 資源的嚴重爭奪。在這種情況下,程序之間會頻繁切換,這增加了程序切換開銷並降低了整體系統效率。
參考 example 倉庫中 Hogwild 實現的程式碼示例來了解 CPU 過度佔用。
在 CPU 上使用 4 個程序執行以下訓練示例命令時
python main.py --num-processes 4
假設機器上有 N 個 vCPU 可用,執行上述命令將生成 4 個子程序。每個子程序會為自己分配 N 個 vCPU,總共需要 4*N 個 vCPU。然而,機器上只有 N 個 vCPU 可用。因此,不同程序會競爭資源,導致頻繁的程序切換。
以下觀察結果表明存在 CPU 過度佔用
高 CPU 利用率:透過使用
htop命令,你可以觀察到 CPU 利用率持續很高,經常達到或超過其最大容量。這表明對 CPU 資源的需求超過了可用的物理核心數,導致程序之間爭奪 CPU 時間。頻繁的上下文切換和低系統效率:在 CPU 過度佔用的情況下,程序爭奪 CPU 時間,作業系統需要快速在不同程序之間切換以公平分配資源。這種頻繁的上下文切換增加了開銷,降低了整體系統效率。
避免 CPU 過度佔用¶
避免 CPU 過度佔用的一個好方法是進行適當的資源分配。確保同時執行的程序或執行緒數量不超過可用的 CPU 資源。
在這種情況下,一個解決方案是在子程序中指定適當的執行緒數。這可以透過在子程序中使用 torch.set_num_threads(int) 函式來設定每個程序的執行緒數實現。
假設機器上有 N 個 vCPU,將生成 M 個程序,則每個程序使用的最大 num_threads 值為 floor(N/M)。為了避免 mnist_hogwild 示例中的 CPU 過度佔用,需要對 example 倉庫中的 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 倍。