Python PyTorch Pipe用法及代碼示例

本文簡要介紹python語言中 torch.distributed.pipeline.sync.Pipe 的用法。

用法:

class torch.distributed.pipeline.sync.Pipe(module, chunks=1, checkpoint='except_last', deferred_batch_norm=False)

參數：

• module(torch.nn.Sequential) -使用流水線並行化的順序模塊。序列中的每個模塊都必須在單個設備上具有所有參數。序列中的每個模塊必須是 nn.Module 或 nn.Sequential (在單個設備上組合多個順序模塊)

• chunks(int) -micro-batches 的數量(默認值：1)

• checkpoint(str) -何時啟用檢查點，'always'、'except_last' 或 'never' 之一(默認值：'except_last')。 'never' 完全禁用檢查點，'except_last' 為除最後一個以外的所有micro-batches 啟用檢查點，'always' 為所有micro-batches 啟用檢查點。

• deferred_batch_norm(bool) -是否使用延遲的BatchNorm移動統計(默認： False )。如果設置為 True ，我們將跟蹤多個 micro-batches 的統計信息，以更新每個 mini-batch 的運行統計信息。

拋出：

• TypeError - 模塊不是 nn.Sequential 。

• ValueError - 無效參數

包裝任意 nn.Sequential 模塊以使用同步管道並行性進行訓練。如果模塊需要大量內存並且不適合單個 GPU，則流水線並行是一種用於訓練的有用技術。

該實現基於torchgpipe 論文。

Pipe 將管道並行性與檢查點相結合，以減少訓練所需的峰值內存，同時最小化設備 under-utilization。

您應該將所有模塊放在適當的設備上並將它們包裝到定義所需執行順序的 nn.Sequential 模塊中。如果模塊不包含任何參數/緩衝區，則假定該模塊應在 CPU 上執行，並且模塊的適當輸入張量在執行前被移動到 CPU。此行為可以被 WithDevice 包裝器覆蓋，該包裝器可用於明確指定模塊應在哪個設備上運行。

例子：

跨 GPU 0 和 1 的兩個 FC 層的管道。

>>> # Need to initialize RPC framework first.
>>> os.environ['MASTER_ADDR'] = 'localhost'
>>> os.environ['MASTER_PORT'] = '29500'
>>> torch.distributed.rpc.init_rpc('worker', rank=0, world_size=1)
>>>
>>> # Build pipe.
>>> fc1 = nn.Linear(16, 8).cuda(0)
>>> fc2 = nn.Linear(8, 4).cuda(1)
>>> model = nn.Sequential(fc1, fc2)
>>> model = Pipe(model, chunks=8)
>>> input = torch.rand(16, 16).cuda(0)
>>> output_rref = model(input)

注意

僅當 Pipe 的檢查點參數為 'never' 時，才可以使用 torch.nn.parallel.DistributedDataParallel 包裝 Pipe 模型。

注意

Pipe 目前僅支持intra-node 流水線，未來將擴展支持inter-node 流水線。轉發函數返回 RRef 以允許將來進行 inter-node 流水線操作，其中輸出可能位於遠程主機上。對於intra-node 流水線，您可以使用 local_value() 在本地檢索輸出。

警告

Pipe 是實驗性的，可能會發生變化。