✨[Feature] Improved parity between convert_method_to_trt_engine
and require_full_compilation=True
#1939
Labels
feature request
New feature or request
Context
Currently in the TorchScript path, the
require_full_compilation=True
applies to conversion operators, or those performing some sort of mathematical operation or transformation. It does not apply to collection packing and unpacking operators. This can lead to confusion in later usage ofconvert_method_to_trt_engine
, since some models which successfully compile withrequire_full_compilation=True
, cannot be converted to fully-TRT engine objects withconvert_method_to_trt_engine
.Proposed Solution
The above is not fully a problem in and of itself, since a model which outputs the schema
(Tensor, (Tensor, Tensor))
could not be converted into a TRT-engine since its output is not a flat list of Tensors, but if its logic could be accelerated to a large extent by Torch-TensorRT, then this could reasonably be called "full compilation". Still, there are certain cases where the parity betweenconvert_method_to_trt_engine
andrequire_full_compilation
could be improved - namely, the case of(Tensor, Tensor, ...)
Tuple-type outputs. These cases can be addressed by removing the extraneousprim::TupleConstruct
calls, and simply returning theIValues
directly from the TRT engine.**Note: ** A recent change to Torch/TorchScript prompted the above problem, as discussed in #1368, by changing the default behavior of traced multi-output-functions from
TensorList
toTuple
.Example
On the above graph,
torch
currently by default considers thec, d
output to be aTuple
, thereby prompting the insertion of aprim::TupleConstruct
operator which is unnecessary. We would prefer if the following (lowered) graph instead were to have bothIValues
registered as outputs, so the entire graph could be converted to a TRT engine viatorch_tensorrt.convert_method_to_trt_engine
.Additional Context
Related to #1938
The text was updated successfully, but these errors were encountered: