Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Extractor dfs performance #1655

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

kylesayrs
Copy link
Contributor

@kylesayrs kylesayrs commented Jun 25, 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

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

@kylesayrs
Copy link
Contributor Author

Related: onnx/onnx#6213

@kylesayrs kylesayrs changed the title WIP: Extractor dfs performance Extractor dfs performance Jul 8, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant