-
Couldn't load subscription status.
- Fork 457
Fix Edge Case: Shared Layer Output Between Model Output and Internal Layers #2407
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
Conversation
...nd also the output of the model.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would not be surprised if some optimizers introduce Identity nodes which won't be removed anymore with this change.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I changed order of optimizer to make my code works.
Does other optimizer make Identity node?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don't remember but since it is a frequent trick developpers use, I would not be surprised. If you have two nodes doing the same thing with the same input, it is simple to remove one node and to replace with an Identity node. They the algorithm removing the identity nodes will do the renaming.
@gramalingam
gramalingam
Sep 26, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I suppose one possibility is to have another copy of IdentityOptimizer at the end?
Uh oh!
There was an error while loading. Please reload this page.
Fix Edge Case: Shared Layer Output Between Model Output and Internal Layers
🎯 Summary
This PR addresses a specific edge case that occurs when converting TensorFlow models to ONNX. It improves the handling of cases where a layer's output is used simultaneously as both the input to another layer and the model's final output.
🔧 Changes
Edge case handling: Resolves the case where a layer's output is used as both the model output and the input to an internal layer.
Improved Transpose output logic:
Uses identity nodes to handle output branching.
Differentiates between model output consumers and internal consumers.
Introduces a two-step process: edge case handling → regular case handling.
Optimizer order optimization: Moves remove_identity to the front of the optimizer chain.
Improved error messages: Provides clearer error messages during graph output validation.
Technical Details
Before:
When a layer's output is connected to both the model output and internal layers, the transpose operation fails, causing the graph structure to break.
After:
Output is branched through identity nodes.
Transpose is applied only to the model output consumer.
Internal consumers maintain the original output.
🔍 Files Changed
tf2onnx/tfonnx.py - Major improvements to the transpose_outputs function.
tf2onnx/graph.py - Improved error messages and code cleanup.
tf2onnx/optimizer/init.py - Changed optimizer order.
Test model file
This file is a test model created by extracting part of the edge case I encountered.
https://drive.google.com/file/d/10ChR5OS4k6az1yG13vdbiassrsxYUkuJ/view?usp=sharing