Add all parameters to from_tensorflow docs (apache#3321)

neo-ai · Jun 27, 2019 · 66ea473 · 66ea473
1 parent 82eaee8
commit 66ea473
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 4 deletions.
diff --git a/nnvm/python/nnvm/frontend/tensorflow.py b/nnvm/python/nnvm/frontend/tensorflow.py
@@ -1188,7 +1188,7 @@ def __init__(self):
         self._input_shapes = {}
 
     def from_tensorflow(self, graph, layout="NHWC", shape=None, outputs=None):
-        """Construct nnvm nodes from tensorflow  graph definition - GraphDef.
+        """Construct nnvm nodes from tensorflow graph definition - GraphDef.
 
         Follow the tensorflow graph definition to parse and convert it to NNVM.
         Some of the assumptions listed below.
@@ -1214,6 +1214,9 @@ def from_tensorflow(self, graph, layout="NHWC", shape=None, outputs=None):
         shape : Dictionary of input dimensions (Optional)
             Graph level input shape dictionary.
 
+        outputs : List of output tensor names (Optional)
+            if not specified then the last node is assumed as graph output.
+
         Returns
         -------
         sym : nnvm.sym.Symbol
@@ -1599,14 +1602,23 @@ def _fix_extranodes(self, op_name, attr, inputs):
         return inputs
 
 def from_tensorflow(graph, layout="NHWC", shape=None, outputs=None):
-    """  Load tensorflow graph which is a python tensorflow graph object into nnvm graph.
+    """Load tensorflow graph which is a python tensorflow graph object into nnvm graph.
     The companion parameters will be handled automatically.
 
     Parameters
     ----------
     graph : GraphDef object
         Tensorflow GraphDef
 
+    layout : target layout to be used (Optional)
+        NCHW only supported now to enable NHWC models on GPU.
+
+    shape : Dictionary of input dimensions (Optional)
+        Graph level input shape dictionary.
+
+    outputs : List of output tensor names (Optional)
+        if not specified then the last node is assumed as graph output.
+
     Returns
     -------
     sym : nnvm.Symbol

diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py
@@ -1787,7 +1787,7 @@ def __init__(self):
         self._branches = {}
 
     def from_tensorflow(self, graph, layout="NHWC", shape=None, outputs=None):
-        """Construct relay nodes from tensorflow  graph definition - GraphDef.
+        """Construct relay nodes from tensorflow graph definition - GraphDef.
 
         Follow the tensorflow graph definition to parse and convert it to Relay.
         Some of the assumptions listed below.
@@ -1813,6 +1813,9 @@ def from_tensorflow(self, graph, layout="NHWC", shape=None, outputs=None):
         shape : Dictionary of input dimensions (Optional)
             Graph level input shape dictionary.
 
+        outputs : List of output tensor names (Optional)
+            if not specified then the last node is assumed as graph output.
+
         Returns
         -------
         sym : relay.op
@@ -2276,14 +2279,23 @@ def _convert_operator(self, op_name, inputs, attrs,
 
 
 def from_tensorflow(graph, layout="NHWC", shape=None, outputs=None):
-    """  Load tensorflow graph which is a python tensorflow graph object into relay.
+    """Load tensorflow graph which is a python tensorflow graph object into relay.
     The companion parameters will be handled automatically.
 
     Parameters
     ----------
     graph : GraphDef object
         Tensorflow GraphDef
 
+    layout : target layout to be used (Optional)
+        NCHW only supported now to enable NHWC models on GPU.
+
+    shape : Dictionary of input dimensions (Optional)
+        Graph level input shape dictionary.
+
+    outputs : List of output tensor names (Optional)
+        if not specified then the last node is assumed as graph output.
+
     Returns
     -------
     sym : relay.op