fix: Bugfix in TRT Engine deserialization indexing by gs-olive · Pull Request #1646 · pytorch/TensorRT (original) (raw)

Description

The IO Tensors and Bindings within the TensorRT ICudaEngine object are not necessarily stored in index-order, or in the order they are stored in PyTorch. Thus, one of the checks in the TRTEngine which extracts the TRT binding name and compares it to the Torch binding name is unnecessary/incorrect and can be improved.

As an example, consider an engine with two inputs: {“input_0”, “input_1”} and two outputs {“output_0”, “output_1”}. The Torch binding names (c10::NameList) stores these as:

[“input_0”, “input_1”, “output_0”, “output_1”]

The TRT Engine binding buffer stores these names as:

[“input_0”, “input_1”, “output_1”, “output_0”]

Thus, when we use direct indexing to access the binding names, an error is encountered, despite the fact that the overall set of IO tensors is the same.

• Fix bug causing crash when loading serialized TRT Engines through Torch-TRT for models with multiple outputs
• Improve TRT Engine binding verification by not assuming sorted indexing in binding order
• Improve check in TRTEngine.cpp for existence of binding index

Fixes #1550
Fixes #1645

Type of change

• Bug fix (non-breaking change which fixes an issue)

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
• [ ~ ] I have added tests to verify my fix or my feature
  • Verified locally on multiple user test cases, listed in the "Fixed" above
• [ 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