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