A GPipe implementation in PyTorch.
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)
What is GPipe?
GPipe is a scalable pipeline parallelism library published by Google Brain, which allows efficient training of large, memory-consuming models. According to the paper, GPipe can train a 25x larger model by using 8x devices (TPU), and train a model 3.5x faster by using 4x devices.
GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
Google trained AmoebaNet-B with 557M parameters over GPipe. This model has achieved 84.3% top-1 and 97.0% top-5 accuracy on ImageNet classification benchmark (the state-of-the-art performance as of May 2019).
Links
Source Code: https://github.com/kakaobrain/torchgpipe
Documentation: https://torchgpipe.readthedocs.io/
Original Paper: https://arxiv.org/abs/1811.06965
Release files for torchgpipe 0.0.7
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| torchgpipe-0.0.7-py3-none-any.whl | Python 3 | none | any | Details |
Release files / torchgpipe-0.0.7-py3-none-any.whl
| Download URL | torchgpipe-0.0.7-py3-none-any.whl |
|---|---|
| Size | 39.1 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
052c704f8c03b7695110e81853e4166c1fc1db3ad1e4cc0fa4eceb8a63e32627
|
|
BLAKE2b-256 checksum How to use checksums |
f48497f3c3b27b666de92477dda6425dd7cb56c7bbaeb115c3a8a1ec7dbe8e05
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/3.2.0 pkginfo/1.5.0.1 requests/2.23.0 setuptools/41.2.0 requests-toolbelt/0.9.1 tqdm/4.49.0 CPython/3.8.1
|