feat: Implement Input class support for FX backend. #1763

peri044 · 2023-03-23T00:49:24Z

Description

This PR adds support for providing torch_tensorrt.Input(shape=x, dtype=x) to FX backend.

Example workflow would be similar to TS backend.

inputs = [torch_tensorrt.Input(shape=(1, 3, 3, 4), dtype=torch.float32)]
trt_mod = torch_tensorrt.compile(
            mod,
            ir="fx",
            inputs=inputs,
            min_acc_module_size=1,
        )

Internally FX performs a) standard lowering passes b) TRT lowering. The standard lowering passes traces through the model which require pytorch tensors for execution.

For standard lowering passes
a) User input of type torch_tensorrt.Input -> Example pytorch tensors (using example_tensors())
b) User input of type torch.Tensors - this type will directly pass through tracing and hence we don't need to do anything.

For TRT lowering
a) self._trt_input is used to convert torch_tensorrt.Input into an InputTensorSpec
b) This can now support dynamic shapes using the same front end interface as Torchscript.

TSInput class is derived from torch_tensorrt.Input because the latter had C++ calls (to_internal). For FX backend, we don't need them for --fx-only installation.
max_batch_size - This has been removed and already deprecated in TRT (not available in TS backend) as well.
FX tracing and lowering previously assumes input tensors reside on the same device as the model parameters. In the case of torch_tensorrt.Input, when we extract example tensors, the default device used is cuda:0. This assumption might not be valid for some non standard use cases. More refactoring is required to ensure we specify the right device and it is being passed accordingly.

Type of change

New feature (non-breaking change which adds functionality)
This change requires a documentation update

Checklist:

My code follows the style guidelines of this project (You can use the linters)
I have performed a self-review of my own code
I have commented my code, particularly in hard-to-understand areas and hacks
I have made corresponding changes to the documentation
I have added tests to verify my fix or my feature
New and existing unit tests pass locally with my changes
I have added the relevant labels to my PR in so that relevant reviewers are notified

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

There are some changes that do not conform to Python style guidelines:

--- py/torch_tensorrt/ts/ts_input.py	2023-03-23 20:35:35.440710 +0000
+++ py/torch_tensorrt/ts/ts_input.py	2023-03-23 20:35:54.340256 +0000
@@ -5,10 +5,11 @@

from torch_tensorrt import _C
from torch_tensorrt import _enums
from torch_tensorrt import _Input
from torch_tensorrt._Input import Input
+

class TSInput(Input):
    """
    Defines an input to a module in terms of expected shape, data type and tensor format.

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to C++ style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: Replace InputTensorSpec with Input Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> feat: Allow torchtrt.Input support for FX backend Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> refactor: Implement conversions from Input -> Pyt tensors, add Input utilities etc. Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: Use InputTensorSpec internally Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: Linter fixes Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: add ts_input.py file Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: Linter fixes Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: minor fixes Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: revert FX changes Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: Address Torchscript test case failures Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: remove device placement of input tensors Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: Linter fixes Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: refactor code Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: Remove max_batch_size and replace generate_input_specs calls Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com> chore: linter fixes Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

github-actions

Code conforms to Python style guidelines

narendasan · 2023-03-29T22:08:18Z

examples/fx/lower_example.py

@@ -98,13 +99,17 @@ def benchmark(

    model = model.cuda().eval()
    inputs = [x.cuda() for x in inputs]
-
+    # inputs = [torch_tensorrt.Input(shape=(128, 3, 224, 224), dtype=torch.float32)]


What are these comments for?

narendasan · 2023-03-29T22:09:12Z

py/torch_tensorrt/_Input.py

@@ -41,6 +40,7 @@ class _ShapeMode(Enum):
    DOMAIN_OFFSET = 2.0
    low_tensor_domain_incl = 0.0
    high_tensor_domain_excl = low_tensor_domain_incl + DOMAIN_OFFSET
+    torch_dtype = None


Should we derive torch_dtype from self.dtype?

narendasan · 2023-03-29T22:09:30Z

py/torch_tensorrt/_Input.py

@@ -173,59 +176,6 @@ def __str__(self) -> str:
        else:
            raise RuntimeError("Unknown input shape mode")

-    def _to_internal(self) -> _C.Input:


Why was this taken out?

narendasan · 2023-03-29T22:10:58Z

py/torch_tensorrt/fx/tools/trt_minimizer.py

    use_experimental_rt: bool = False,
 ) -> TRTModule:
    interp = TRTInterpreter(
-        mod, InputTensorSpec.from_tensors(inputs), explicit_batch_dimension=True
+        mod, _Input.Input.from_tensors(inputs), explicit_batch_dimension=True


Can this just be Input.from_tensors

narendasan · 2023-03-29T22:12:27Z

py/torch_tensorrt/ts/ts_input.py

+from torch_tensorrt._Input import Input
+
+
+class TSInput(Input):


This should be a hidden class, we dont want people using this, handle any conversion to TSInput internally

Yeah that's the intended usage. In the current design, for torchscript, they still use the same interface
torch_tensorrt.Input() and we internally convert into TSInput in the ts/_compile_spec.py file. Anything different that you have in mind ?

No this would be fine

frank-wei

Overall, I can not approve this PR for some reasons:

fx2trt is used internally and core APIs are called by different services. So far, we did not build the Import mechanism to protect the internal service. (ImportIt is used in other Pytorch projects which will run internal CI that covers not only unit tests but also other prod tests for any PR). So we can not take the risk to take any potential big change since that will break our internal services.
Due to the above reason, any change for FX path needs to consider backward compatibility. Also consider to reduce the PR size for easy review.
cc @yinghai @wushirong

frank-wei · 2023-03-30T06:11:39Z

py/torch_tensorrt/fx/fx2trt.py

@@ -153,7 +153,6 @@ def validate_conversion(self):

    def run(
        self,
-        max_batch_size=64,


I am afraid we can not do this change. We have to maintain backward compatibility on the API, otherwise, it will break our internal product.

frank-wei · 2023-03-30T06:13:22Z

py/torch_tensorrt/fx/input_tensor_spec.py

@@ -4,58 +4,59 @@

 from .types import Shape, ShapeRange
 from .utils import get_dynamic_dims
-
-
-def generate_input_specs(inputs, lower_setting, additional_inputs=None):


Same here. These API are used in many internal products.

narendasan · 2023-03-30T23:12:36Z

@frank-wei can we expose the old API somewhere else for backwards compatibility but move to the unified one for users? One of the big problems right now is it is difficult to move from torchscript to fx since many of the settings are named differently or used differently

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to C++ style guidelines

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions

Code conforms to Python style guidelines

github-actions

Code conforms to C++ style guidelines

peri044 · 2023-04-05T01:19:27Z

Closing this in favour of #1807

frank-wei · 2023-04-12T05:33:09Z

@frank-wei can we expose the old API somewhere else for backwards compatibility but move to the unified one for users? One of the big problems right now is it is difficult to move from torchscript to fx since many of the settings are named differently or used differently

Specifically, I am thinking the unified version could happen on aten2trt?

frank-wei · 2023-04-12T05:45:37Z

py/torch_tensorrt/fx/input_tensor_spec.py

@@ -116,6 +117,43 @@ def from_tensors(cls, tensors: Sequence[torch.Tensor]) -> List["InputTensorSpec"
        assert isinstance(tensors, (list, tuple))
        return [cls.from_tensor(t) for t in tensors]

+    @classmethod


I assume we use the input_obj is general interface.

frank-wei · 2023-04-12T06:02:20Z

py/torch_tensorrt/fx/passes/lower_pass_manager_builder.py

@@ -262,7 +268,13 @@ def _default_replace_mutable_op_pass(self) -> PassManager:
    def build_trt_lower_pipeline(
        self, input: Input, additional_input: Optional[Input] = None
    ) -> PassManager:
-        self._input = input
+


can we start a new func instead of change this build_trt_lower_pipeline behavior?

frank-wei · 2023-04-12T06:04:46Z

py/torch_tensorrt/fx/passes/pass_utils.py

@@ -21,6 +22,30 @@
 FINAL_CHECK_RTOL_MULTIPLIER: float = 10


+def extract_example_tensors_from_input(


can we pass real tensor as input to the lowering workflow? Why do we need the input_obj as input?

peri044 added the WIP Work is in progress, pull request should not be merged yet label Mar 23, 2023

peri044 requested a review from narendasan March 23, 2023 00:49

facebook-github-bot added cla signed fx labels Mar 23, 2023

github-actions bot added component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: fx component: runtime labels Mar 23, 2023

github-actions bot approved these changes Mar 23, 2023

View reviewed changes

This comment was marked as resolved.

Sign in to view

github-actions bot approved these changes Mar 23, 2023

View reviewed changes

github-actions bot requested changes Mar 23, 2023

View reviewed changes

github-actions bot approved these changes Mar 24, 2023

View reviewed changes

peri044 mentioned this pull request Mar 27, 2023

✨[Feature] Enable conversion of Pytorch tensors <=> Input class in FX frontend #1718

Closed

peri044 force-pushed the fx_input branch from 987b2b3 to afd725e Compare March 28, 2023 07:10

github-actions bot added the component: tests Issues re: Tests label Mar 28, 2023

github-actions bot approved these changes Mar 28, 2023

View reviewed changes

This comment was marked as resolved.

Sign in to view

github-actions bot approved these changes Mar 28, 2023

View reviewed changes

This comment was marked as resolved.

Sign in to view

github-actions bot approved these changes Mar 28, 2023

View reviewed changes

peri044 force-pushed the fx_input branch from 2f3ada3 to 4f07f58 Compare March 29, 2023 21:07

github-actions bot approved these changes Mar 29, 2023

View reviewed changes

This comment was marked as resolved.

Sign in to view

peri044 removed the WIP Work is in progress, pull request should not be merged yet label Mar 29, 2023

peri044 force-pushed the fx_input branch from 4f07f58 to e5c2026 Compare March 29, 2023 21:27

github-actions bot approved these changes Mar 29, 2023

View reviewed changes

peri044 requested a review from frank-wei March 29, 2023 21:35

github-actions bot approved these changes Mar 29, 2023

View reviewed changes

narendasan reviewed Mar 29, 2023

View reviewed changes

frank-wei reviewed Mar 30, 2023

View reviewed changes

chore: address review comments

792b231

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Mar 30, 2023

View reviewed changes

This comment was marked as spam.

Sign in to view

chore: fix comma

0c6a27d

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

This comment was marked as off-topic.

Sign in to view

github-actions bot approved these changes Mar 31, 2023

View reviewed changes

chore: Linter fixes

78eb3b5

Signed-off-by: Dheeraj Peri <peri.dheeraj@gmail.com>

github-actions bot approved these changes Mar 31, 2023

View reviewed changes

peri044 closed this Apr 5, 2023

frank-wei reviewed Apr 12, 2023

View reviewed changes

		from torch_tensorrt._Input import Input


		class TSInput(Input):

		@@ -21,6 +22,30 @@
		FINAL_CHECK_RTOL_MULTIPLIER: float = 10


		def extract_example_tensors_from_input(

feat: Implement Input class support for FX backend. #1763

feat: Implement Input class support for FX backend. #1763

Conversation

peri044 commented Mar 23, 2023 • edited Loading

Description

Type of change

Checklist:

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

This comment was marked as resolved.

This comment was marked as resolved.

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

This comment was marked as resolved.

github-actions bot left a comment

Choose a reason for hiding this comment

This comment was marked as resolved.

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

This comment was marked as resolved.

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

frank-wei left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

narendasan commented Mar 30, 2023

github-actions bot left a comment

Choose a reason for hiding this comment

This comment was marked as spam.

This comment was marked as off-topic.

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

peri044 commented Apr 5, 2023

frank-wei commented Apr 12, 2023

frank-wei Apr 12, 2023 • edited Loading

peri044 commented Mar 23, 2023 •

edited

Loading

frank-wei left a comment •

edited

Loading

frank-wei Apr 12, 2023 •

edited

Loading