Skip to content

Commit

Permalink
Create onnxvis entry point; fix when main graph does not have a name (
Browse files Browse the repository at this point in the history
#22)

Also fixed a bug when the main graph does not have a name
  • Loading branch information
justinchuby authored May 19, 2024
1 parent dceb257 commit 6be2167
Show file tree
Hide file tree
Showing 6 changed files with 44 additions and 5 deletions.
4 changes: 4 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -158,3 +158,7 @@ cython_debug/
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/

*.onnx
*.textproto
*.pb
7 changes: 6 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,12 @@ pip install --upgrade model-explorer-onnx

```bash
model-explorer --extensions=model_explorer_onnx

# Or as a shortcut
onnxvis

# Supply model path
onnxvis model.onnx
```

## Screenshots
Expand All @@ -29,4 +35,3 @@ model-explorer --extensions=model_explorer_onnx
<img width="1293" alt="image" src="https://github.com/justinchuby/model-explorer-onnx/assets/11205048/fbf2fa05-bd29-4938-93d1-709690d9f9c6">

<img width="1301" alt="image" src="https://github.com/justinchuby/model-explorer-onnx/assets/11205048/a68f7ecd-1fa1-4eac-9e1f-8e9a5bbf9fe3">

9 changes: 6 additions & 3 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,10 @@ description = "Adapter for ai-edge-model-explorer to support ONNX models"
authors = [{ name = "Justin Chu", email = "justinchu@microsoft.com" }]
readme = "README.md"
requires-python = ">=3.8"
license = { file = "LICENSE" }
license = {text = "MIT License"}
keywords = ["onnx", "model-explorer", "visualization"]
classifiers = [
"Development Status :: 3 - Alpha",
"Environment :: Console",
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"Operating System :: POSIX",
"Operating System :: MacOS :: MacOS X",
Expand All @@ -33,5 +33,8 @@ dependencies = [
"ml_dtypes",
]

[project.scripts]
onnxvis = "model_explorer_onnx.bin.onnxvis:main"

[project.urls]
Repository = "https://github.com/justinchuby/model-explorer-onnx"
Empty file.
21 changes: 21 additions & 0 deletions src/model_explorer_onnx/bin/onnxvis.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
#!/usr/bin/env python
"""A shortcut to run model explorer with ONNX extension."""

import subprocess
import sys


def main():
# Run model explorer
subprocess.run(
[
"model-explorer",
"--extensions",
"model_explorer_onnx",
*sys.argv[1:],
]
)


if __name__ == "__main__":
main()
8 changes: 7 additions & 1 deletion src/model_explorer_onnx/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -482,10 +482,16 @@ def convert(
if opset_version is None:
opset_version = _DEFAULT_OPSET_VERSION
# TODO: Better support subgraphs in nodes
if model.graph.name is None:
model.graph.name = "<main>"
logger.warning(
"Main graph of ONNX file '%s' does not have a name. Set name to '<main>'",
model_path,
)
main_graph = create_graph(
model.graph, all_function_ids, opset_version=opset_version
)
assert main_graph is not None
assert main_graph is not None, "Bug: Main graph should not be None"
graphs.append(main_graph)

for function in model.functions.values():
Expand Down

0 comments on commit 6be2167

Please sign in to comment.