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倍的设备