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}
}

展开

附加信息