Prepare API docs

docs/Makefile

@@ -0,0 +1,19 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line.
+SPHINXOPTS    =
+SPHINXBUILD   = sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

docs/conf.py

@@ -0,0 +1,69 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# http://www.sphinx-doc.org/en/master/config
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+import os
+import sys
+sys.path.insert(0, os.path.abspath('..'))
+
+
+# -- Project information -----------------------------------------------------
+
+project = 'torchgpipe'
+copyright = '2019, Kakao Brain'
+author = 'Kakao Brain'
+
+# The full version, including alpha/beta/rc tags
+release = __import__('torchgpipe').__version__
+
+
+# -- General configuration ---------------------------------------------------
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    # We follow Google style docstrings just like PyTorch.
+    'sphinx.ext.napoleon',
+]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ['_templates']
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+
+# Link to PyTorch's documentation.
+extensions.append('sphinx.ext.intersphinx')
+intersphinx_mapping = {'torch': ('https://pytorch.org/docs/stable/', None)}
+
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = 'alabaster'
+
+html_theme_options = {
+    'description': 'GPipe in PyTorch',
+
+    'github_user': 'KakaoBrain',
+    'github_repo': 'torchgpipe',
+    'github_type': 'star',
+}
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ['_static']

docs/index.rst

@@ -0,0 +1,48 @@
+torchgpipe
+==========
+
+A GPipe_ implementation in PyTorch_.
+
+.. _GPipe: https://arxiv.org/abs/1811.06965
+.. _PyTorch: https://pytorch.org/
+
+.. sourcecode:: python
+
+   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:
+       outputs = model(input)
+
+
+Installing
+----------
+
+torchgpipe is available on PyPI_.
+
+.. sourcecode:: console
+
+   $ pip install torchgpipe
+
+.. _PyPI: https://pypi.org/project/torchgpipe
+
+
+API
+---
+
+.. autoclass:: torchgpipe.GPipe(module, balance, \**kwargs)
+
+   .. automethod:: forward(input)
+
+
+Licensing and Authors
+---------------------
+
+This package is opened under the Apache License 2.0.
+
+We are `Heungsub Lee`_ and Myungryong Jeong in `Kakao Brain`_.
+
+.. _Heungsub Lee: https://subl.ee/
+.. _Kakao Brain: https://kakaobrain.com/

docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=.
+set BUILDDIR=_build
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.http://sphinx-doc.org/
+	exit /b 1
+)
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS%
+
+:end
+popd

torchgpipe/gpipe.py

@@ -41,48 +41,40 @@ def __lt__(self, other: tuple) -> bool:
 
 
 class GPipe(nn.Module):
-    """Implements `GPipe <https://arxiv.org/abs/1811.06965>`_ a scalable
-    pipeline parallelism at the module level. GPipe aims to reduce peak memory
-    usage while minimizing device under-untilization.
-
-    This container separates a sequential module into several partitions over
-    each device just like naive model parallelism. Its :meth:`forward` divides
-    a mini-batch into multiple micro-batches. Each partition sends a
-    micro-batch output to the next partition then takes the next micro-batch as
-    soon as possible. So unlike naive model parallelism, those partitions can
-    be overlapped with each other.
-
-    GPipe also integrates checkpointing to futher reduce memory usage.
-    Recomputation of a micro-batch can be evaluated preceding its backward
-    pass.
-
-    .. note::
-        There are still idle time called "bubbling" so the utilization might be
-        less than 100%. But the bubbling can be partly shrunk by increasing the
-        number of micro-batches or balanced partitioning.
+    """Wraps an arbitrary :class:`~torch.nn.Sequential` module to train on
+    GPipe_. If the module requires lots of memory, GPipe will be very
+    efficient::
+
+        model = torch.nn.Sequential(l1, l2, l3, l4)
+        gpipe = GPipe(model, balance=[1,1,1,1], chunks=8)
+        output = gpipe(input)
+
+    .. _GPipe: https://arxiv.org/abs/1811.06965
+
+    GPipe combines pipeline parallelism with checkpointing to reduce peak
+    memory required to train while minimizing device under-utilization.
+
+    You should determine the balance when defining a GPipe module, as balancing
+    will not be done automatically. The module will be partitioned into
+    multiple devices according to the given balance. You may rely on heuristics
+    to find your own optimal configuration.
 
     Args:
         module (nn.Sequential):
             sequential module to be parallelized
         balance (ints):
-            lengths of layers on each partition
+            list of number of layers in each partition
 
     Keyword Args:
         devices (iterable of devices):
             devices to use (default: all CUDA devices)
         chunks (int):
             number of micro-batches (default: 1)
         checkpoint (str):
-            when to enable checkpoints, one of 'always', 'except_last', or
+            when to enable checkpointing, one of 'always', 'except_last', or
             'never' (default: 'except_last')
         deferred_batch_norm (bool):
-            enable deferred BatchNorm moving statistics (default: False)
-
-    Example::
-
-        >>> model = torch.nn.Sequential(l1, l2, l3, l4)
-        >>> gpipe = GPipe(model, balance=[1, 1, 1, 1], chunks=8)
-        >>> output = gpipe(input)
+            whether to use deferred BatchNorm moving statistics (default: False)
 
     """
 
@@ -351,6 +343,19 @@ def pull_output(self,
         return output
 
     def forward(self, input: TensorOrTensors) -> TensorOrTensors:
+        """:class:`GPipe` is a fairly transparent module wrapper. It doesn't
+        modify the input and output signature of the underlying module. But
+        there's type restriction. Input and output have to be a
+        :class:`~torch.Tensor` or a tuple of tensors. This restriction is
+        applied at partition boundaries too.
+
+        Args:
+            input (tensor or tensors): input mini-batch
+
+        Returns:
+            tensor or tensors: output mini-batch
+
+        """
         in_queue, out_queue = self.spawn_workers()
         num_inputs = self.push_input(input, in_queue)
         return self.pull_output(num_inputs, in_queue, out_queue)