Use sets to improve performance of dfs search #6213
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
5714577 to
61e7fb1
Compare
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #6213 +/- ##
==========================================
+ Coverage 56.95% 57.11% +0.15%
==========================================
Files 506 506
Lines 30467 31087 +620
Branches 4592 4624 +32
==========================================
+ Hits 17353 17754 +401
- Misses 12285 12505 +220
+ Partials 829 828 -1 ☔ View full report in Codecov by Sentry. |
c1ca823 to
e75c67a
Compare
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
e75c67a to
01d56d4
Compare
|
It seems reachable + unreachable = set of nodes. Do you it is doable to remoave unreachable and simplify the function signature assuming performance stays the same? |
|
@xadupre The nodes_to_search = [
index
for index in range(len(nodes))
if node_output_name in nodes[index].output and index not in reachable
]This leads to a longer runtime for larger models
|
xadupre
approved these changes
Jul 5, 2024
andife
pushed a commit
to andife/onnx
that referenced
this pull request
Jul 20, 2024
### Description
This change modifies the DFS search used in model extraction to use sets
rather than iterables
### Motivation and Context
While shape inference is the most significant bottleneck, these changes
are a step in the direction of being able to support model extraction
for very large graphs.
## Test Script ##
```python3
import onnx
from onnx.utils import Extractor
model = onnx.load("obertquant.onnx")
extractor = Extractor(model)
extracted_model = extractor.extract_model(
input_names=["input_ids", "attention_mask", "token_type_ids"], output_names=["2058"]
)
onnx.save(extracted_model, "truncated.onnx")
```
Benchmarks were produced using pyinstrument and analyzing the
`Extractor.extract_model` function
| Model Name | Num Nodes | Previous | New |
| ---- | -------- | --------- | ----- |
| obertquant.onnx | 1271 | 0.158s | 0.110s |
| ai-town-3B.onnx | 3515 | 8.002s | 3.725s |
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Andreas Fehlner <fehlner@arcor.de>
linshokaku
pushed a commit
to linshokaku/onnx
that referenced
this pull request
Oct 2, 2024
### Description
This change modifies the DFS search used in model extraction to use sets
rather than iterables
### Motivation and Context
While shape inference is the most significant bottleneck, these changes
are a step in the direction of being able to support model extraction
for very large graphs.
## Test Script ##
```python3
import onnx
from onnx.utils import Extractor
model = onnx.load("obertquant.onnx")
extractor = Extractor(model)
extracted_model = extractor.extract_model(
input_names=["input_ids", "attention_mask", "token_type_ids"], output_names=["2058"]
)
onnx.save(extracted_model, "truncated.onnx")
```
Benchmarks were produced using pyinstrument and analyzing the
`Extractor.extract_model` function
| Model Name | Num Nodes | Previous | New |
| ---- | -------- | --------- | ----- |
| obertquant.onnx | 1271 | 0.158s | 0.110s |
| ai-town-3B.onnx | 3515 | 8.002s | 3.725s |
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Linsho Kaku <linsho@preferred.jp>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This change modifies the DFS search used in model extraction to use sets rather than iterables
Motivation and Context
While shape inference is the most significant bottleneck, these changes are a step in the direction of being able to support model extraction for very large graphs.
Test Script
Benchmarks were produced using pyinstrument and analyzing the
Extractor.extract_modelfunction