"graphbolt/src/cuda/index_select_csc_impl.hip" did not exist on "3795a006b91c94291b911f0daa261c0598d7ffd8"
Unverified Commit 20278c0d authored by msbaines's avatar msbaines Committed by GitHub
Browse files

[docs] add a Pipeline tutorial (#82)

parent 61aac92a
[settings]
known_third_party =numpy,pytest,pytorch_sphinx_theme,recommonmark,setuptools,torch,torchtext,torchvision
known_third_party =numpy,pytest,recommonmark,setuptools,torch,torchtext,torchvision
API Reference
=============
.. toctree::
:maxdepth: 1
optim/oss
nn/pipe
......@@ -7,16 +7,7 @@ Welcome to fairscale's documentation!
=====================================
.. toctree::
:maxdepth: 1
:caption: API Reference
:maxdepth: 2
api/optim/oss
api/nn/pipe
Indices and tables
==================
* :ref:`genindex`
* :ref:`modindex`
* :ref:`search`
api/index
tutorials/index
Tutorials
=========
.. toctree::
:maxdepth: 1
pipe
Pipeline Parallel
=================
Let us start with a toy model that contains two linear layers.
.. code-block:: default
import torch
import torch.nn as nn
class ToyModel(nn.Module):
def __init__(self):
super(ToyModel, self).__init__()
self.net1 = torch.nn.Linear(10, 10)
self.relu = torch.nn.ReLU()
self.net2 = torch.nn.Linear(10, 5)
def forward(self, x):
x = self.relu(self.net1(x))
return self.net2(x)
model = ToyModel()
To run this model on 2 GPUs we need to convert the model
to ``torch.nn.Sequential`` and then wrap it with ``fairscale.nn.Pipe``.
.. code-block:: default
import fairscale
import torch
import torch.nn as nn
model = nn.Sequential(
torch.nn.Linear(10, 10),
torch.nn.ReLU(),
torch.nn.Linear(10, 5)
)
model = fairscale.nn.Pipe(model, balance=[2, 1])
This will run the first two layers on ``cuda:0`` and the last
layer on ``cuda:1``. To learn more, visit the `Pipe <../api/nn/pipe.html>`_ documentation.
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment