Add example script (static to dynamic batch size)

Signed-off-by: Rajeev Rao <[email protected]>

tools/onnx-graphsurgeon/examples/10_dynamic_batch_size/README.md

@@ -0,0 +1,40 @@
+# Dynamic Batch Size
+
+## Introduction
+
+This example first generates a basic model with static batch size in the input,
+then modifies the resulting model with dynamic batch size.
+
+## Basics: modifying the input's batch size
+Below is the main code necessary to convert a static ONNX model to dynamic:
+```python
+graph = gs.import_onnx(onnx_model)
+for input in graph.inputs:
+   input.shape[0] = 'N'
+```
+
+The above code should be enough for simple models.
+However, some models might also need their internal layers updated, such as static `Reshape` layers.
+Please see the section below for an example of this scenario. 
+
+## Running the example
+
+1. Generate a model with static input shape and several nodes (including static `Reshape`), and save it to `model.onnx` by running:
+    ```bash
+    python3 generate.py
+    ```
+
+    ![../resources/10_model.onnx.png](../resources/10_model.onnx.png)
+
+2. Convert the model to support dynamic batch size, and save it to `modified.onnx` by running:
+    ```bash
+    python3 modify.py
+    ```
+
+    This script does the following:
+    - Updates the inputs' batch size to `N` (dynamic symbol)
+    - Updates the shape of `Reshape` nodes (if they exist) to `-1` (dynamic value)
+    
+    The resulting graph now supports any batch size:
+
+    ![../resources/10_dynamic.onnx.png](../resources/10_dynamic.onnx.png)

tools/onnx-graphsurgeon/examples/10_dynamic_batch_size/generate.py

@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import onnx_graphsurgeon as gs
+import numpy as np
+import onnx
+
+
+##########################################################################################################
+# Register functions to simplify the graph building process later on.
+
+@gs.Graph.register()
+def conv(self, inp, weights, dilations, group, strides):
+    out = self.layer(
+        op="Conv",
+        inputs=[inp, weights],
+        outputs=["conv_out"],
+        attrs={"dilations": dilations, "group": group, "kernel_shape": weights.shape[2:], "strides": strides},
+    )[0]
+    out.dtype = inp.dtype
+    return out
+
+
+@gs.Graph.register()
+def reshape(self, data, shape):
+    out = self.layer(op="Reshape", inputs=[data, shape], outputs=["reshape_out"])[0]
+    out.dtype = data.dtype
+    return out
+
+
+@gs.Graph.register()
+def matmul(self, lhs, rhs):
+    out = self.layer(op="MatMul", inputs=[lhs, rhs], outputs=["matmul_out"])[0]
+    out.dtype = lhs.dtype
+    return out
+
+##########################################################################################################
+
+
+# Set input
+X = gs.Variable(name="input_1", dtype=np.float32, shape=(1, 3, 28, 28))
+graph = gs.Graph(inputs=[X])
+
+# Connect intermediate tensors
+conv_out = graph.conv(
+    X, weights=np.ones(shape=(32, 3, 3, 3), dtype=np.float32), dilations=[1, 1], group=1, strides=[1, 1]
+)
+reshape_out = graph.reshape(conv_out, np.array([1, 21632], dtype=np.int64))
+matmul_out = graph.matmul(reshape_out, np.ones(shape=(21632, 10), dtype=np.float32))
+
+# Set output
+graph.outputs = [matmul_out]
+
+# Save graph
+onnx.save(gs.export_onnx(graph), "model.onnx")

tools/onnx-graphsurgeon/examples/10_dynamic_batch_size/modify.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import onnx
+import onnx_graphsurgeon as gs
+
+# Load ONNX model
+graph = gs.import_onnx(onnx.load("model.onnx"))
+
+# Update input shape
+for input in graph.inputs:
+    input.shape[0] = 'N'
+
+# Update 'Reshape' nodes (if they exist)
+reshape_nodes = [node for node in graph.nodes if node.op == "Reshape"]
+for node in reshape_nodes:
+    # The batch dimension in the input shape is hard-coded to a static value in the original model.
+    # To make the model work with our dynamic batch size, we can use a `-1`, which indicates that the
+    # dimension should be automatically determined.
+    node.inputs[1].values[0] = -1
+
+# Save dynamic model
+onnx.save(gs.export_onnx(graph), "dynamic.onnx")

tools/onnx-graphsurgeon/tests/test_examples.py

@@ -48,6 +48,7 @@ def __init__(self, name, infer=True):
     ("07_creating_a_model_with_the_layer_api", [Artifact("model.onnx")]),
     ("08_replacing_a_subgraph", [Artifact("model.onnx"), Artifact("replaced.onnx")]),
     ("09_shape_operations_with_the_layer_api", [Artifact("model.onnx")]),
+    ("10_dynamic_batch_size", [Artifact("model.onnx"), Artifact("dynamic.onnx")]),
 ]
 
 # Extract any ``` blocks from the README