fix: Allow full model compilation with collection inputs (input_signature) by gs-olive · Pull Request #1656 · pytorch/TensorRT (original) (raw)
Description
- Allow users to specify full model compilation when using
input_signature, which allows for complex collection-based inputs - Enable "psuedo-partitioning" phase for input collections as well as output collections
- Update
OutputIsCollectionto include dictionary outputs, and add functionInputIsCollectionto detect collection-based inputs during graph compilation - Remove automatic fallback for collection pack/unpack operations when using
input_signatureargument - Add collections tests to ensure full compilation is respected for input and output collections
Fixes #1602
Fixes #1293
Addresses #1595
Type of change
- Bug fix
- New feature
Checklist:
- [ x ] My code follows the style guidelines of this project (You can use the linters)
- [ x ] I have performed a self-review of my own code
- [ x ] I have commented my code, particularly in hard-to-understand areas and hacks
- [ x ] I have made corresponding changes to the documentation
- [ x ] I have added tests to verify my fix or my feature
- [ x ] New and existing unit tests pass locally with my changes
- [ x ] I have added the relevant labels to my PR in so that relevant reviewers are notified
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
- Allow users to specify full model compilation when using
input_signature, which allows for complex collection-based inputs - Enable "psuedo-partitioning" phase for input collections as well as output collections
- Update
OutputIsCollectionto include dictionary outputs, and add functionInputIsCollectionto detect collection-based inputs during graph compilation - Remove automatic fallback for collection pack/unpack operations when
using
input_signatureargument - Add collections tests to ensure full compilation is respected for input and output collections
- Add documentation to
READMEfor usage of input signature - Add documentation to "Getting Started" page for usage of input signature
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM