Merge branch 'main'

Author: Wei Wei

fx2trt/.gitignore

@@ -0,0 +1,6 @@
+build/
+dist/
+*__pycache__*
+.DS_Store
+fx2trt_oss.egg-info/
+

fx2trt/LICENSE

@@ -0,0 +1,28 @@
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in the
+   documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors
+   may be used to endorse or promote products derived from this software
+   without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

fx2trt/README.md

@@ -0,0 +1,29 @@
+# FX2TRT
+This package provide pure eager-mode tooling to convert a PyTorch nn.Module to a TensorRT engine.
+
+## Installation
+First, let's install PyTorch.
+```
+conda install -y pytorch cudatoolkit=11.3 -c pytorch-nightly
+```
+Then, you need to install your TensorRT and it's pythin binding
+```
+tar -xzvf TensorRT-8.2.1.8.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz
+export LD_LIBRARY_PATH=$HOME/TensorRT-8.2.1.8/lib:$HOME/TensorRT-8.2.1.8/targets/x86_64-linux-gnu/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+cd TensorRT-8.2.1.8/python
+python3 -m pip install tensorrt-8.2.1.8-cp36-none-linux_x86_64.whl
+```
+Then, it's simply as this.
+```
+cd fx2trt
+python setup.py install
+```
+## Test
+Follow instruction in [pytorch/benchmark](https://github.com/pytorch/benchmark) to setup some benchmarks.
+
+Then try using `--fx2trt` for individual cases, e.g.
+```
+cd benchmark
+python run.py resnet50 -d cuda -t eval -m eager --fx2trt
+```
+And you should expect TensorRT logs being printed and the case ran through without error.

fx2trt/fx/__init__.py

@@ -0,0 +1,10 @@
+from .converters import *  # noqa: F403 F401
+from .converter_registry import (  # noqa
+    CONVERTERS,
+    NO_EXPLICIT_BATCH_DIM_SUPPORT,
+    NO_IMPLICIT_BATCH_DIM_SUPPORT,
+    tensorrt_converter,
+)
+from .fx2trt import TRTInterpreter, TRTInterpreterResult  # noqa
+from .input_tensor_spec import InputTensorSpec  # noqa
+from .trt_module import TRTModule  # noqa

fx2trt/fx/converter_registry.py

@@ -0,0 +1,31 @@
+from typing import Any, Callable, Dict
+
+from torch.fx.node import Target
+
+
+CONVERTERS: Dict[Target, Any] = {}
+NO_IMPLICIT_BATCH_DIM_SUPPORT = {}
+NO_EXPLICIT_BATCH_DIM_SUPPORT = {}
+
+
+def tensorrt_converter(
+    key: Target,
+    no_implicit_batch_dim: bool = False,
+    no_explicit_batch_dim: bool = False,
+    enabled: bool = True,
+) -> Callable[[Any], Any]:
+    def register_converter(converter):
+        CONVERTERS[key] = converter
+        if no_implicit_batch_dim:
+            NO_IMPLICIT_BATCH_DIM_SUPPORT[key] = converter
+        if no_explicit_batch_dim:
+            NO_EXPLICIT_BATCH_DIM_SUPPORT[key] = converter
+        return converter
+
+    def disable_converter(converter):
+        return converter
+
+    if enabled:
+        return register_converter
+    else:
+        return disable_converter

fx2trt/fx/converters/__init__.py

@@ -0,0 +1,18 @@
+# @manual=//deeplearning/trt/python:py_tensorrt
+import tensorrt as trt
+
+if hasattr(trt, "__version__"):
+    from .activation import *  # noqa: F401 F403
+    from .adaptive_avgpool import *  # noqa: F401 F403
+    from .add import *  # noqa: F401 F403
+    from .batchnorm import *  # noqa: F401 F403
+    from .convolution import *  # noqa: F401 F403
+    from .linear import *  # noqa: F401 F403
+    from .maxpool import *  # noqa: F401 F403
+    from .mul import *  # noqa: F401 F403
+    from .transformation import *  # noqa: F401 F403
+    from .quantization import *  # noqa: F401 F403
+    from .acc_ops_converters import *  # noqa: F401 F403
+
+    TRT_LOGGER = trt.Logger()
+    trt.init_libnvinfer_plugins(TRT_LOGGER, "")