torchgpipe下載torchgpipe源代碼下載

torchgpipe

Python

v0.0.7

下載

Torchgpipe

Pytorch中的GPIPE實現。它是針對CUDA而不是TPU優化的。

 from torchgpipe import GPipe
model = nn . Sequential ( a , b , c , d )
model = GPipe ( model , balance = [ 1 , 1 , 1 , 1 ], chunks = 8 )
output = model ( input )

什麼是gpipe？

GPIPE是由Google Brain發表的可擴展管道並行庫，可有效培訓大型內存耗盡模型。根據論文，GPIPE可以使用8倍設備（TPU）訓練25倍較大的型號，並使用4倍設備訓練3.5倍的型號。

GPIPE：使用管道並行性對巨型神經網絡進行有效培訓

Google培訓了Amoebanet-B，gpipe的參數為5.57億。該模型在ImageNet分類基準（截至2019年5月的最新性能）上獲得了84.3％的TOP-1和97.0％的前5位準確性。

GPIPE使用（a）管道並行性和（b）向後傳播過程中正向傳播的自動重建，因此利用訓練大型模型。在Pytorch社區中著名的術語之後，我們將（b）稱為檢查點。

管道並行性: GPIPE將模型分為多個分區，並將每個分區放在其他設備上以佔據更多的內存容量。它將迷你批量分成多個微批次，以使分區盡可能平行。
檢查點: 檢查點應用於每個分區，以最大程度地減少模型的整體內存消耗。在向前傳播過程中，只記住了分區之間邊界的張量。所有其他中間張量均揮發，並在必要時在返回過程中重新計算。

用法

目前，Torchgpipe需要以下環境：

Python 3.6+
Pytorch 1.1+

要使用Torchgpipe，請通過PYPI安裝它：

$ pip install torchgpipe

要用GPIPE訓練模塊，只需用torchgpipe.GPipe將其包裝。您的模塊必須為nn.Sequential隨著GPIPE的序列，將自動將模塊分為連續層的分區。 balance參數確定每個分區中的層數。 chunks參數指定了微批次的數量。輸入，輸出和中間張量必須是Tensor或Tuple[Tensor, ...] 。

下面的示例代碼顯示瞭如何將四層分為單層的四個分區分為四個分區。此代碼還將一個小批量分成8個微批次：

 from torchgpipe import GPipe

model = nn . Sequential ( a , b , c , d )
model = GPipe ( model , balance = [ 1 , 1 , 1 , 1 ], chunks = 8 )

for input in data_loader :
    output = model ( input )

文件

有關更多信息，包括API參考，請訪問Torchgpipe.ReadThedocs.io。

基準測試

torchgpipe.readthedocs.io中提供了全部詳細信息和更多基準。

RESNET-101精度基準

批量大小	Torchgpipe	nn.dataparallel	Goyal等人。
256	21.99±0.13	22.02±0.11	22.08±0.06
1k	22.24±0.19	22.04±0.24	N/A。
4K	22.13±0.09	N/A。	N/A。

GPIPE應該是透明的，不要引入其他高參數調整。為了驗證透明度，我們在Imagenet上重現了RESNET-101的TOP-1錯誤率，如Goyal等人的精確，大型Minibatch SGD的表2（C）所報導。

U-NET（B，C）內存基準測試

實驗	U-NET（B，C）	參數	內存使用
基線	（6，72）	3622m	20.3吉布
管道1	（11，128）	2.21b	20.5 Gib
管道2	（24，128）	4.99b	43.4 gib
管道4	（24，160）	7.80b	79.1吉布
管道8	（48，160）	15.82b	154.1 gib

該表顯示了GPIPE如何促進縮放U-NET模型。基線表示沒有管道並行的基線或檢查點，而管道-1 ， -2 ， -4 ， -8表示模型是用GPIPE訓練的，該模型具有相應數量的分區。

在這裡，我們使用了簡化的U-NET體系結構。模型的大小由分別與層和過濾器數量成正比的超參數B和C確定。

U-NET（5，64）速度基準

實驗	吞吐量	加快
基線	28.500/s	1×
管道1	24.456/s	0.858×
管道2	35.502/s	1.246×
管道4	67.042/s	2.352×
管道8	88.497/s	3.105×

為了通過跳過連接來驗證效率，我們用不同數量的設備測量了U-NET的吞吐量。我們選擇使用U-NET，因為它具有多個長跳連接。

Amoebanet-D（18，256）速度基準

實驗	吞吐量	Torchgpipe	Huang等。
n = 2，m = 1	26.733/s	1×	1×
n = 2，m = 4	41.133/s	1.546×	1.07×
n = 2，m = 32	47.386/s	1.780×	1.21×
n = 4，m = 1	26.827/s	1.006×	1.13×
n = 4，m = 4	44.543/s	1.680×	1.26×
n = 4，m = 32	72.412/s	2.711×	1.84×
n = 8，m = 1	24.918/s	0.932×	1.38×
n = 8，m = 4	70.065/s	2.625×	1.72×
n = 8，m = 32	132.413/s	4.966×	3.48×

（ n ：分區數， M ：微批量的數量）

該表顯示了Amoebanet-D（18，256）上的速度基準，如Huang等人的Gpipe表2所報導。請注意，我們用n替換了紙。

筆記

該項目是功能性的，但是尚未確認接口。所有公共API都會發生變化，而無需警告，直到v0.1.0。

作者和許可

Torchgpipe項目由Heungsub Lee，Myungryong Jeong和Kakao Brain的Chiheon Kim開發，並與Sungbin Lim，Ildoo Kim，Woonhyuk Baek和Boogeon Yoon的幫助。它是根據BSD許可證的第3條規定的。

引用

如果您將此庫應用於任何項目和研究，請引用我們的代碼：

 @article{kim2020torchgpipe,
    title={torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models},
    author={Chiheon Kim and Heungsub Lee and Myungryong Jeong and Woonhyuk Baek and Boogeon Yoon and Ildoo Kim and Sungbin Lim and Sungwoong Kim},
    year={2020},
    eprint={2004.09910},
    archivePrefix={arXiv}
}

展開

附加信息