Improve model Extractor #6920
Improve model Extractor #6920
Conversation
Improve Extractor to handle the case where user requests a graph output, that is simultaneously a intermediate value, to be extracted as an input of the new graph. Previously the value info of the graph i/o are not added to the mapping, causing _collect_new_io_core to fail with key error Signed-off-by: Justin Chu <justinchu@microsoft.com>
There was a problem hiding this comment.
Pull Request Overview
This PR improves the Extractor to correctly map intermediate graph I/O values when a graph output that is also an intermediate value is requested, preventing KeyError failures.
- Removed the previous vimap and wmap in favor of initializers and value_infos, which are updated with input and output values.
- Consolidated the collection of new I/O into a single method (_collect_new_io) and updated tensor collection accordingly.
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
92e4c22 to
d00ffa0
Compare
There was a problem hiding this comment.
Pull Request Overview
This PR improves the Extractor in onnx/utils.py to correctly handle cases in which a graph output (acting as an intermediate value) is also requested as an input, preventing a KeyError.
- Replaces the separate _collect_new_inputs/_outputs methods with a unified _collect_new_io method that validates io names against value_infos.
- Updates mapping references from vimap/wmap to value_infos/initializers throughout the extraction process.
Comments suppressed due to low confidence (1)
onnx/utils.py:42
- [nitpick] Consider renaming _collect_new_io to a more descriptive name, such as _collect_valid_io, to clearly indicate that it validates and collects tensor info for both inputs and outputs.
def _collect_new_io(self, io_names_to_extract: list[str]) -> list[ValueInfoProto]:
Codecov ReportAttention: Patch coverage is
✅ All tests successful. No failed tests found.
Additional details and impacted files@@ Coverage Diff @@
## main #6920 +/- ##
==========================================
- Coverage 56.31% 56.30% -0.02%
==========================================
Files 509 509
Lines 32590 32582 -8
Branches 3097 3095 -2
==========================================
- Hits 18354 18346 -8
Misses 13378 13378
Partials 858 858 ☔ View full report in Codecov by Sentry. |
zhenhuaw-me
left a comment
There was a problem hiding this comment.
LGTM. Better to have some tolerance about the ONNX spec.
FYI: our TensorRT project has a improved version (and many more powerful features) baked in polygraphy: https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/examples/cli/surgeon/01_isolating_subgraphs
Improve Extractor to handle the case where user requests a graph output, that is simultaneously an intermediate value, to be extracted as an input of the new graph. Previously the value info of the graph i/o are not added to the mapping, causing _collect_new_io_core to fail with KeyError. --------- Signed-off-by: Justin Chu <justinchu@microsoft.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: gramalingam <grama@microsoft.com>
Improve Extractor to handle the case where user requests a graph output, that is simultaneously an intermediate value, to be extracted as an input of the new graph.
Previously the value info of the graph i/o are not added to the mapping, causing _collect_new_io_core to fail with KeyError.