From d3c98b91980182faef25772f061fe7adedb14262 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Fri, 26 Apr 2024 01:24:19 +0000 Subject: [PATCH 1/3] Tutorial --- docs/examples/ir/tutorial.ipynb | 326 ++++++++++++++++++++++++++++++++ 1 file changed, 326 insertions(+) create mode 100644 docs/examples/ir/tutorial.ipynb diff --git a/docs/examples/ir/tutorial.ipynb b/docs/examples/ir/tutorial.ipynb new file mode 100644 index 000000000..4e4ed9ed9 --- /dev/null +++ b/docs/examples/ir/tutorial.ipynb @@ -0,0 +1,326 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Intro to the ONNX IR\n", + "\n", + "https://github.com/microsoft/onnxscript/tree/main/onnxscript/ir\n", + "\n", + "The ONNX IR is our attempt at providing a robust, efficient and Pythonic in-memory IR for ONNX to power model building, analysis and manipulation. It has\n", + "\n", + "- **Full ONNX spec support**: all valid models representable by ONNX protobuf, and a subset of invalid models (so you can load and fix them).\n", + "- **Low memory footprint**: mmap'ed external tensors; unified interface for ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size limitation. Zero copies.\n", + "- **Straightforward access patterns**: Access value information and traverse the graph topology at ease.\n", + "- **Robust mutation support**: Create as many iterators as you like on the graph while mutating it.\n", + "- **Speed**: Performant graph manipulation, serialization/deserialization to Protobuf.\n", + "- **Pythonic and familiar APIs**: Classes define Pythonic apis and still map to ONNX protobuf concepts in an intuitive way.\n", + "- **No protobuf dependency**: The IR does not require protobuf once the model is converted to the IR representation, decoupling from the serialization format." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import onnx\n", + "from onnxscript import ir" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Load ONNX model\n", + "\n", + "proto = onnx.load(\"/home/justinchu/dev/onnx-script/testdata/e2e_models/mobilenetv2_100/dynamo/mobilenetv2_100_dynamo.onnx\")\n", + "model = ir.serde.deserialize_model(proto)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph = model.graph" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.display()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(graph.initializers.keys())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.initializers[\"conv_stem.weight\"].display()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# graph.initializers[\"model.embed_tokens.weight\"].numpy()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "len(graph)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "node = graph[6]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(node)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "new_node = ir.Node(\n", + " \"my_custom_domain\",\n", + " \"Linear_classifier\",\n", + " node.inputs,\n", + " name=\"new_torch_nn_modules_linear_Linear_classifier_1\"\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "new_node.display()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "for value, replacement in zip(node.outputs, new_node.outputs):\n", + " for user_node, index in tuple(value.uses()):\n", + " user_node.replace_input_with(index, replacement)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.insert_after(node, new_node)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(graph)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.remove(node, safe=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.outputs[0] = new_node.outputs[0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(graph)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(node.inputs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.remove(node, safe=True)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "new_model_proto = ir.serde.serialize_model(model)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.outputs[0].shape = ir.Shape([1, 1000])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.outputs[0].dtype = ir.DataType.FLOAT" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.outputs[0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(graph)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "new_model_proto = ir.serde.serialize_model(model)\n", + "onnx.save(new_model_proto, \"new_model_proto.onnx\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "len(node.inputs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "len(node.outputs)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "graph.insert_after" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "onnx", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.9" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} From 817a53db0027e4a8fb3092eb13ce24aa44eaea07 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 30 Apr 2024 01:32:07 +0000 Subject: [PATCH 2/3] Create intro --- docs/examples/ir/tutorial.ipynb | 326 ------------------ .../introduction.py | 127 +++++++ 2 files changed, 127 insertions(+), 326 deletions(-) delete mode 100644 docs/examples/ir/tutorial.ipynb create mode 100644 docs/intermediate_representation/introduction.py diff --git a/docs/examples/ir/tutorial.ipynb b/docs/examples/ir/tutorial.ipynb deleted file mode 100644 index 4e4ed9ed9..000000000 --- a/docs/examples/ir/tutorial.ipynb +++ /dev/null @@ -1,326 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Intro to the ONNX IR\n", - "\n", - "https://github.com/microsoft/onnxscript/tree/main/onnxscript/ir\n", - "\n", - "The ONNX IR is our attempt at providing a robust, efficient and Pythonic in-memory IR for ONNX to power model building, analysis and manipulation. It has\n", - "\n", - "- **Full ONNX spec support**: all valid models representable by ONNX protobuf, and a subset of invalid models (so you can load and fix them).\n", - "- **Low memory footprint**: mmap'ed external tensors; unified interface for ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size limitation. Zero copies.\n", - "- **Straightforward access patterns**: Access value information and traverse the graph topology at ease.\n", - "- **Robust mutation support**: Create as many iterators as you like on the graph while mutating it.\n", - "- **Speed**: Performant graph manipulation, serialization/deserialization to Protobuf.\n", - "- **Pythonic and familiar APIs**: Classes define Pythonic apis and still map to ONNX protobuf concepts in an intuitive way.\n", - "- **No protobuf dependency**: The IR does not require protobuf once the model is converted to the IR representation, decoupling from the serialization format." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import onnx\n", - "from onnxscript import ir" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Load ONNX model\n", - "\n", - "proto = onnx.load(\"/home/justinchu/dev/onnx-script/testdata/e2e_models/mobilenetv2_100/dynamo/mobilenetv2_100_dynamo.onnx\")\n", - "model = ir.serde.deserialize_model(proto)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "model" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph = model.graph" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.display()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(graph.initializers.keys())" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.initializers[\"conv_stem.weight\"].display()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# graph.initializers[\"model.embed_tokens.weight\"].numpy()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "len(graph)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "node = graph[6]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(node)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "new_node = ir.Node(\n", - " \"my_custom_domain\",\n", - " \"Linear_classifier\",\n", - " node.inputs,\n", - " name=\"new_torch_nn_modules_linear_Linear_classifier_1\"\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "new_node.display()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "for value, replacement in zip(node.outputs, new_node.outputs):\n", - " for user_node, index in tuple(value.uses()):\n", - " user_node.replace_input_with(index, replacement)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.insert_after(node, new_node)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(graph)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.remove(node, safe=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.outputs[0] = new_node.outputs[0]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(graph)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(node.inputs)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.remove(node, safe=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "new_model_proto = ir.serde.serialize_model(model)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.outputs[0].shape = ir.Shape([1, 1000])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.outputs[0].dtype = ir.DataType.FLOAT" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.outputs[0]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "print(graph)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "new_model_proto = ir.serde.serialize_model(model)\n", - "onnx.save(new_model_proto, \"new_model_proto.onnx\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "len(node.inputs)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "len(node.outputs)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "graph.insert_after" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "onnx", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.9" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/docs/intermediate_representation/introduction.py b/docs/intermediate_representation/introduction.py new file mode 100644 index 000000000..228bf70a4 --- /dev/null +++ b/docs/intermediate_representation/introduction.py @@ -0,0 +1,127 @@ +# %% [markdown] +# # ONNX IR Introduction +# ## Overview +# +# The ONNX IR a robust, efficient and Pythonic in-memory IR for ONNX to power +# model building, analysis and manipulation. It has +# +# - **Full ONNX spec support**: all valid models representable by ONNX protobuf, +# and a subset of invalid models (so you can load and fix them). +# - **Low memory footprint**: mmap'ed external tensors; unified interface for +# ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size limitation. +# Zero copies. +# - **Straightforward access patterns**: Access value information and traverse the +# graph topology at ease. +# - **Robust mutation support**: Create as many iterators as you like on the graph +# while mutating it. +# - **Speed**: Performant graph manipulation, serialization/deserialization to Protobuf. +# - **Pythonic and familiar APIs**: Classes define Pythonic apis and still map +# to ONNX protobuf concepts in an intuitive way. +# - **No protobuf dependency**: The IR does not require protobuf once the model +# is converted to the IR representation, decoupling from the serialization +# format. +# +# This tutorial will demonstrate how you can use the ONNX IR to inspect, manipulate +# and build ONNX models. + +# %% +import onnx +from onnxscript import ir + +# %% +# Load ONNX model + +proto = onnx.load("/home/justinchu/dev/onnx-script/testdata/e2e_models/mobilenetv2_100/dynamo/mobilenetv2_100_dynamo.onnx") +model = ir.serde.deserialize_model(proto) + +# %% +model + +# %% +graph = model.graph + +# %% +graph.display() + +# %% +print(graph.initializers.keys()) + +# %% +graph.initializers["conv_stem.weight"].display() + +# %% +# graph.initializers["model.embed_tokens.weight"].numpy() + +# %% +len(graph) + +# %% +node = graph[6] + +# %% +print(node) + +# %% +new_node = ir.Node( + "my_custom_domain", + "Linear_classifier", + node.inputs, + name="new_torch_nn_modules_linear_Linear_classifier_1" +) + +# %% +new_node.display() + +# %% +for value, replacement in zip(node.outputs, new_node.outputs): + for user_node, index in tuple(value.uses()): + user_node.replace_input_with(index, replacement) + +# %% +graph.insert_after(node, new_node) + +# %% +print(graph) + +# %% +graph.remove(node, safe=True) + +# %% +graph.outputs[0] = new_node.outputs[0] + +# %% +print(graph) + +# %% +print(node.inputs) + +# %% +graph.remove(node, safe=True) + +# %% +new_model_proto = ir.serde.serialize_model(model) + +# %% +graph.outputs[0].shape = ir.Shape([1, 1000]) + +# %% +graph.outputs[0].dtype = ir.DataType.FLOAT + +# %% +graph.outputs[0] + +# %% +print(graph) + +# %% +new_model_proto = ir.serde.serialize_model(model) +onnx.save(new_model_proto, "new_model_proto.onnx") + +# %% +len(node.inputs) + +# %% +len(node.outputs) + +# %% +graph.insert_after From be5dd1c9d249f8ebb4d2c12ca93738fc91006860 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 30 Apr 2024 18:36:32 +0000 Subject: [PATCH 3/3] intro --- docs/intermediate_representation/introduction.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/docs/intermediate_representation/introduction.py b/docs/intermediate_representation/introduction.py index 228bf70a4..1bb174b30 100644 --- a/docs/intermediate_representation/introduction.py +++ b/docs/intermediate_representation/introduction.py @@ -23,11 +23,24 @@ # # This tutorial will demonstrate how you can use the ONNX IR to inspect, manipulate # and build ONNX models. - +# +# ## Installation +# +# The ONNX IR is part of the `onnxscript` package. You can install it via pip: +# ```bash +# pip install onnxscript +# ``` +# Then import the package: +# # %% import onnx + from onnxscript import ir +# %% [markdown] +# ## Data Structures +# The ONNX IR models an ONNX graph similar to the ONNX protobuf definitions. You +# expect familiar classes like `ir.Model`, `ir.Graph`, `ir.Node`, and `ir.Tensor`. # %% # Load ONNX model