Add auto doc gen for ORTModule API during CI build (#7046)

In addition to ORTModule auto documentation during packaging, this PR also update golden numbers to fix CI

.gitignore

@@ -24,16 +24,16 @@ nuget_root/
 *.code-workspace
 __pycache__
 onnxruntime_profile*.json
-/docs/python/*.md
-/docs/python/auto_examples/*
-/docs/python/media/*
-/docs/python/examples/*.onnx
-/docs/python/examples/graph.*
+/docs/python/inference/*.md
+/docs/python/inference/auto_examples/*
+/docs/python/inference/media/*
+/docs/python/inference/examples/*.onnx
+/docs/python/inference/examples/graph.*
 /docs/python/*_LICENSE
 /csharp/**/obj/
 /csharp/**/bin/
 /csharp/Directory.Build.props
-docs/python/*.onnx
+docs/python/inference/*.onnx
 *.onnx
 onnxprofile_profile_test_*.json
 /csharp/packages

docs/python/inference/conf.py

@@ -30,7 +30,6 @@ extensions = [
     "sphinx.ext.autodoc",
     'sphinx.ext.githubpages',
     "sphinx_gallery.gen_gallery",
-    'sphinx.ext.autodoc',
     'sphinx.ext.graphviz',
     "pyquickhelper.sphinxext.sphinx_runpython_extension",
 ]

docs/python/inference/examples/plot_backend.py

@@ -20,10 +20,16 @@ from onnxruntime.capi.onnxruntime_pybind11_state import InvalidArgument
 import onnxruntime.backend as backend
 from onnx import load
 
+########################################
+# The device depends on how the package was compiled,
+# GPU or CPU.
+from onnxruntime import get_device
+device = get_device()
+
 name = datasets.get_example("logreg_iris.onnx")
 model = load(name)
 
-rep = backend.prepare(model, 'CPU')
+rep = backend.prepare(model, device)
 x = np.array([[-1.0, -2.0]], dtype=np.float32)
 try:
     label, proba = rep.run(x)
@@ -32,17 +38,11 @@ try:
 except (RuntimeError, InvalidArgument) as e:
     print(e)
 
-########################################
-# The device depends on how the package was compiled,
-# GPU or CPU.
-from onnxruntime import get_device
-print(get_device())
-
 ########################################
 # The backend can also directly load the model
 # without using *onnx*.
 
-rep = backend.prepare(name, 'CPU')
+rep = backend.prepare(name, device)
 x = np.array([[-1.0, -2.0]], dtype=np.float32)
 try:
     label, proba = rep.run(x)

docs/python/requirements.txt

@@ -4,7 +4,10 @@ scikit-learn
 skl2onnx
 sphinx
 sphinx-gallery
-pyquickhelper
 sphinxcontrib.imagesvg
+sphinx_rtd_theme
+pyquickhelper
 tensorflow
 tf2onnx
+pandas
+pydot

docs/python/training/conf.py

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+# -*- coding: utf-8 -*-
+#
+# Configuration file for the Sphinx documentation builder.
+
+# -- Project information -----------------------------------------------------
+
+project = 'ORTModule'
+copyright = '2018-2021, Microsoft'
+author = 'Microsoft'
+version = '0.1' # TODO: Should use `onnxruntime.__version__` instead?
+release = version
+
+# -- General configuration ---------------------------------------------------
+
+extensions = ['sphinx.ext.autodoc',
+              'sphinx.ext.intersphinx'
+]
+templates_path = ['_templates']
+exclude_patterns = []
+
+# -- Options for HTML output -------------------------------------------------
+
+html_theme = 'sphinx_rtd_theme'
+html_static_path = ['_static']
+
+# -- Options for intersphinx extension ---------------------------------------
+
+intersphinx_mapping = {'https://docs.python.org/': None}

docs/python/training/content.rst

@@ -0,0 +1,54 @@
+This document describes ORTModule PyTorch frontend API for the ONNX Runtime (aka ORT) training acelerator.
+
+What is new
+===========
+
+Version 0.1
+-----------
+
+#. Initial version
+
+Overview
+========
+The aim of ORTModule is to provide a drop-in replacement for one or more torch.nn.Module objects in a user’s PyTorch program,
+and execute the forward and backward passes of those modules using ORT.
+
+As a result, the user will be able to accelerate their training script gradually using ORT,
+without having to modify their training loop.
+
+Users will be able to use standard PyTorch debugging techniques for convergence issues,
+e.g. by probing the computed gradients on the model’s parameters.
+
+The following code example illustrates how ORTModule would be used in a user’s training script,
+in the simple case where the entire model can be offloaded to ONNX Runtime:
+
+.. code-block:: python
+
+    # Original PyTorch model
+    class NeuralNet(torch.nn.Module):
+        def __init__(self, input_size, hidden_size, num_classes):
+            ...
+        def forward(self, x): 
+            ...
+
+    model = NeuralNet(input_size=784, hidden_size=500, num_classes=10)
+    model = ORTModule(model) # Only change to original PyTorch script
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
+
+    # Training Loop is unchanged
+    for data, target in data_loader:
+        optimizer.zero_grad()
+        output = model(data)
+        loss = criterion(output, target)
+        loss.backward()
+        optimizer.step()
+
+API
+===
+
+.. automodule:: onnxruntime.training.ortmodule
+    :members:
+    :show-inheritance:
+    :member-order: bysource
+..    :inherited-members:

docs/python/training/index.rst

@@ -0,0 +1,24 @@
+.. ORTModule documentation master file, created by
+   sphinx-quickstart on Thu Mar 11 10:21:00 2021.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+ORTModule documentation
+=======================
+
+ORTModule is a PyTorch frontend API for `ONNX Runtime <https://www.onnxruntime.ai/docs/>`_,
+a cross-platform inferencing and training accelerator.
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+   content
+
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`

orttraining/orttraining/test/python/orttraining_test_orttrainer_frontend.py

@@ -1493,14 +1493,14 @@ def _adam_max_norm_clip_data():
         ]
     elif device_capability_major == 5:  # M60 for CI machines (Python Packaging Pipeline)
         return [
-            (0, 'cuda', 1.0, 1, 12, [10.564176,  9.97518 ,  9.474583,  9.076513,  8.724233,  8.440754,\
-                    8.156107,  7.905789,  7.683374,  7.420595,  7.156939,  6.914139]),
-            (0, 'cuda', 0.1, 1, 12, [10.564176,  9.975904,  9.475933,  9.078633,  8.727232,  8.444615,\
-                    8.160932,  7.911339,  7.689524,  7.427308,  7.164138,  6.921518]),
-            (42, 'cuda', 1.0, 1, 12, [10.660578, 10.027208,  9.518457,  9.10457 ,  8.767458,  8.469093,\
-                    8.207001,  7.92541 ,  7.655277,  7.434964,  7.155968,  6.924634]),
-            (42, 'cuda', 0.1, 1, 12, [10.660578, 10.027987,  9.519927,  9.106762,  8.770505,  8.473034,\
-                    8.211944,  7.931111,  7.661622,  7.441899,  7.163355,  6.932114]),
+            (0, 'cuda', 1.0, 1, 12, [10.618382, 10.08292 ,  9.603334,  9.258133,  8.917768,  8.591574,
+                                     8.318401,  8.042292,  7.783608,  7.50226 ,  7.236041,  7.035602]),
+            (0, 'cuda', 0.1, 1, 12, [10.618382, 10.083632,  9.604639,  9.260109,  8.920504,  8.595082,
+                                     8.322799,  8.047493,  7.78929 ,  7.508382,  7.242587,  7.042367]),
+            (42, 'cuda', 1.0, 1, 12, [10.68639 , 10.102986,  9.647681,  9.293091,  8.958928,  8.625297,
+                                      8.351107,  8.079577,  7.840723,  7.543044,  7.284141,  7.072688]),
+            (42, 'cuda', 0.1, 1, 12, [10.68639 , 10.103672,  9.649025,  9.295167,  8.961777,  8.629059,
+                                      8.355571,  8.084871,  7.846589,  7.549438,  7.290722,  7.079446]),
         ]
 @pytest.mark.parametrize("seed,device,max_norm_clip,gradient_accumulation_steps,total_steps,expected_loss", _adam_max_norm_clip_data())
 def testORTTrainerAdamMaxNormClip(seed, device, max_norm_clip, gradient_accumulation_steps, total_steps, expected_loss):
@@ -1542,14 +1542,14 @@ def _lamb_max_norm_clip_data():
         ]
     elif device_capability_major == 5:  # M60 for CI machines (Python Packaging Pipeline)
         return [
-            (0, 'cuda', 1.0, 1, 12, [10.564176, 10.429815, 10.331507, 10.261825, 10.19336 , 10.110986,\
-                    10.041771,  9.990074,  9.985901,  9.892414,  9.819457,  9.753627]),
-            (0, 'cuda', 0.1, 1, 12, [10.564176, 10.391491, 10.253088, 10.146585, 10.044328,  9.930671,\
-                    9.830513,  9.752279,  9.72234 ,  9.606323,  9.506898,  9.417118]),
-            (42, 'cuda', 1.0, 1, 12, [10.660578, 10.510471, 10.431763, 10.358577, 10.301462, 10.209934,\
-                    10.167318, 10.03529 ,  9.995482,  9.938999,  9.875689,  9.80955 ]),
-            (42, 'cuda', 0.1, 1, 12, [10.660578, 10.471846, 10.352203, 10.241209, 10.149426, 10.026606,\
-                    9.952093,  9.792846,  9.726216,  9.645785,  9.556379,  9.467741]),
+            (0, 'cuda', 1.0, 1, 12, [10.618382, 10.50222 , 10.403347, 10.35298 , 10.288447, 10.237399,
+                                     10.184225, 10.089048, 10.008952,  9.972644,  9.897674,  9.84524 ]),
+            (0, 'cuda', 0.1, 1, 12, [10.618382, 10.466732, 10.330871, 10.24715 , 10.150972, 10.069127,
+                                     9.98974 ,  9.870169,  9.763693,  9.704323,  9.605957,  9.533117]),
+            (42, 'cuda', 1.0, 1, 12, [10.68639 , 10.511692, 10.447308, 10.405255, 10.334866, 10.261473,
+                                      10.169422, 10.107138, 10.069889,  9.97798 ,  9.928105,  9.896435]),
+            (42, 'cuda', 0.1, 1, 12, [10.68639 , 10.477489, 10.376671, 10.301725, 10.200718, 10.098477,
+                                      9.97995 ,  9.890104,  9.828899,  9.713555,  9.639567,  9.589856]),
         ]
 @pytest.mark.parametrize("seed,device,max_norm_clip, gradient_accumulation_steps,total_steps,expected_loss", _lamb_max_norm_clip_data())
 def testORTTrainerLambMaxNormClip(seed, device, max_norm_clip, gradient_accumulation_steps, total_steps, expected_loss):

requirements-doc.txt

@@ -1,3 +1,3 @@
 sphinx
 sphinx_gallery
-
+sphinx_rtd_theme

tools/ci_build/build.py

@@ -86,7 +86,7 @@ def _openvino_verify_device_type(device_read):
         comma_separated_devices = device_read.split(":")
         comma_separated_devices = comma_separated_devices[1].split(',')
         if (len(comma_separated_devices) < 2):
-            print("Atleast two devices required in Hetero Mode")
+            print("At least two devices required in Hetero Mode")
             status_hetero = False
         dev_options = ["CPU", "GPU", "MYRIAD", "FPGA", "HDDL"]
         for dev in comma_separated_devices:
@@ -183,11 +183,15 @@ def parse_arguments():
         help="""When running the Test phase, run symbolic shape inference against
         available test data directories.""")
 
-    # generate documentaiton
+    # generate documentation
     parser.add_argument(
         "--gen_doc", action='store_true',
         help="Generate documentation on contrib ops")
 
+    parser.add_argument(
+        "--gen-api-doc", action='store_true',
+        help="Generate API documentation for PyTorch frontend")
+
     # CUDA related
     parser.add_argument("--use_cuda", action='store_true', help="Enable CUDA.")
     parser.add_argument(
@@ -1783,6 +1787,9 @@ def main():
     if args.code_coverage and not args.android:
         raise BuildError("Using --code_coverage requires --android")
 
+    if args.gen_api_doc and len(args.config) != 1:
+        raise BuildError('Using --get-api-doc requires a single build config')
+
     # Disabling unit tests for VAD-F as FPGA only supports
     # models with NCHW layout
     if args.use_openvino == "VAD-F_FP32":
@@ -1995,6 +2002,12 @@ def main():
     if args.gen_doc and (args.build or args.test):
         generate_documentation(source_dir, build_dir, configs)
 
+    if args.gen_api_doc and (args.build or args.test):
+        print('Generating Python doc for ORTModule...')
+        docbuild_dir = os.path.join(source_dir, 'tools', 'doc')
+        run_subprocess(['bash', 'builddoc.sh', os.path.dirname(sys.executable),
+                        source_dir, build_dir, args.config[0]], cwd=docbuild_dir)
+
     log.info("Build complete")
 
 

tools/ci_build/github/azure-pipelines/templates/py-packaging-stage.yml

@@ -122,16 +122,16 @@ stages:
               -e NIGHTLY_BUILD \
               -e BUILD_BUILDNUMBER \
               onnxruntimecpubuild \
-                bash /onnxruntime_src/tools/doc/builddoc.sh $(PythonManylinuxDir)/bin/ /onnxruntime_src /build
+                bash /onnxruntime_src/tools/doc/builddoc.sh $(PythonManylinuxDir)/bin/ /onnxruntime_src /build Release
           workingDirectory: $(Build.SourcesDirectory)
 
       - task: CopyFiles@2
         displayName: 'Copy Python Documentation to: $(Build.ArtifactStagingDirectory)'
         condition: ne(variables['PythonVersion'], '3.9')  # tensorflow not available on python 3.9
         inputs:
-          SourceFolder: '$(Build.BinariesDirectory)/docs/html'
+          SourceFolder: '$(Build.BinariesDirectory)/docs/inference/html'
           Contents: '**'
-          TargetFolder: '$(Build.ArtifactStagingDirectory)'
+          TargetFolder: '$(Build.ArtifactStagingDirectory)/inference_html_doc'
 
       - task: PublishBuildArtifacts@1
         displayName: 'Publish Artifact: ONNXRuntime python wheel and documentation'
@@ -284,8 +284,34 @@ stages:
           Contents: 'Release/dist/*.whl'
           TargetFolder: '$(Build.ArtifactStagingDirectory)'
 
+      - task: CmdLine@2
+        displayName: 'Build Python Documentation'
+        condition: ne(variables['PythonVersion'], '3.9')  # tensorflow not available on python 3.9
+        inputs:
+          script: |
+            mkdir -p $HOME/.onnx
+            docker run --rm \
+              --volume /data/onnx:/data/onnx:ro \
+              --volume $(Build.SourcesDirectory):/onnxruntime_src \
+              --volume $(Build.BinariesDirectory):/build \
+              --volume /data/models:/build/models:ro \
+              --volume $HOME/.onnx:/home/onnxruntimedev/.onnx \
+              -e NIGHTLY_BUILD \
+              -e BUILD_BUILDNUMBER \
+              onnxruntimetraininggpubuild \
+                bash /onnxruntime_src/tools/doc/builddoc.sh $(PythonManylinuxDir)/bin/ /onnxruntime_src /build Release
+          workingDirectory: $(Build.SourcesDirectory)
+
+      - task: CopyFiles@2
+        displayName: 'Copy Python Documentation to: $(Build.ArtifactStagingDirectory)'
+        condition: ne(variables['PythonVersion'], '3.9')  # tensorflow not available on python 3.9
+        inputs:
+          SourceFolder: '$(Build.BinariesDirectory)/docs/training/html'
+          Contents: '**'
+          TargetFolder: '$(Build.ArtifactStagingDirectory)/training_html_doc'
+
       - task: PublishBuildArtifacts@1
-        displayName: 'Publish Artifact: ONNXRuntime python wheel'
+        displayName: 'Publish Artifact: ONNXRuntime python wheel and documentation'
         inputs:
           ArtifactName: onnxruntime_gpu
 

tools/ci_build/github/linux/docker/manylinux2014_build_scripts/build.sh

@@ -87,6 +87,7 @@ yum -y install \
     ${TOOLCHAIN_DEPS} \
     diffutils \
     gettext \
+    graphviz \
     file \
     kernel-devel \
     libffi-devel \

tools/ci_build/github/linux/docker/scripts/manylinux/install_deps.sh

@@ -112,4 +112,3 @@ make install
 
 cd /
 rm -rf /tmp/src
-

tools/doc/builddoc.sh

@@ -1,16 +1,25 @@
 # This script must be executed from this folder.
+
 # $1 python path
 # $2 source folder
 # $3 build folder
-echo "----"
-echo $1
-echo $2
-echo $3
-echo "----"
-ls $3/Release
-echo "----"
+# $4 build config
+
+# Install doc generation tools
 $1/python -m pip install -r $2/docs/python/requirements.txt
-export PYTHONPATH=$3/Release:$PYTHONPATH
-$1/python -m sphinx -j1 -v -T -b html -d $3/docs/_doctrees/html $2/docs/python $3/docs/html
-$1/python -u $2/tools/doc/rename_folders.py $3/docs/html
-# zip -r $3/python_doc.zip $3/docs/html
+
+# Fake onnxruntime installation
+export PYTHONPATH=$3/$4:$PYTHONPATH
+
+# Remove old docs
+rm -rf $3/docs/
+
+# Inference doc
+$1/python -m sphinx -j1 -v -T -b html -d $3/docs/inference/_doctrees/html $2/docs/python/inference $3/docs/inference/html
+$1/python -u $2/tools/doc/rename_folders.py $3/docs/inference/html
+# (cd $3/docs/inference/html && zip -r $3/docs/python_inference_doc.zip .)
+
+# Training doc
+$1/python -m sphinx -j1 -v -T -b html -d $3/docs/training/_doctrees/html $2/docs/python/training $3/docs/training/html
+$1/python -u $2/tools/doc/rename_folders.py $3/docs/training/html
+# (cd $3/docs/training/html && zip -r $3/docs/python_training_doc.zip .)