文件名称:torchgpipe:PyTorch中的GPipe实现
文件大小:182KB
文件格式:ZIP
更新时间:2024-02-25 17:25:02
deep-learning pytorch parallelism model-parallelism 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倍的设备