Merge branch 'master' of https://github.com/NVIDIA/TRTorch

Author: narendasan

.github/ISSUE_TEMPLATE/op-converter-request.md

@@ -0,0 +1,22 @@
+---
+name: "\U00002194 Op Converter Request"
+about: Submit a proposal/request to support a new PyTorch operator in TRTorch
+
+---
+
+##
+<!-- Name of operator to add support for -->
+
+- **Function Schema**:
+
+- **Original PyTorch API**:
+
+- **Relevant TensorRT Documentation**:
+
+## Alternatives
+
+<!-- A clear and concise description of any alternative solutions you've considered, if any. -->
+
+## Additional context
+
+<!-- Add any other context or screenshots about the operator request here. -->

BUILD

@@ -11,6 +11,7 @@ pkg_tar(
         "//core/conversion/var:include",
         "//core/conversion/tensorcontainer:include",
         "//core/conversion/evaluators:include",
+        "//core/conversion/converters/impl/plugins:include",
         "//core/execution:include",
         "//core/lowering:include",
         "//core/lowering/passes:include",
@@ -47,7 +48,6 @@ pkg_tar(
 )
 
 
-
 pkg_tar(
     name = "libtrtorch",
     extension = "tar.gz",

README.md

@@ -203,10 +203,6 @@ Thanks for wanting to contribute! There are two main ways to handle supporting a
 
 You can register a converter for your op using the `NodeConverterRegistry` inside your application.
 
-## Known Limitations
-
-- You cannot use both Adaptive Pooling in PyTorch and also use TRTorch Dynamic input shape (follow [#49](https://github.com/NVIDIA/TRTorch/issues/49) for the latest on the issue)
-
 ## Structure of the repo
 
 | Component     | Description                                                  |

core/conversion/converters/BUILD

@@ -10,7 +10,7 @@ config_setting(
 cc_library(
     name = "converters",
     hdrs = [
-        "converters.h",
+        "converters.h"
     ],
     srcs = [
         "NodeConverterRegistry.cpp",
@@ -28,14 +28,17 @@ cc_library(
         "impl/shuffle.cpp",
         "impl/softmax.cpp",
         "impl/unary.cpp",
-        "impl/interpolate.cpp"
+        "impl/interpolate.cpp",
+        "impl/select.cpp",
+        "impl/stack.cpp"
     ],
     deps = [
         "@tensorrt//:nvinfer",
         "//core/util:prelude",
         "//core/conversion/var",
         "//core/conversion/tensorcontainer",
         "//core/conversion/conversionctx",
+        "//core/conversion/converters/impl/plugins"
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default":  ["@libtorch//:libtorch"],

core/conversion/converters/impl/interpolate.cpp

@@ -1,8 +1,9 @@
 #include "torch/torch.h"
 #include "core/util/prelude.h"
 #include "core/conversion/converters/converters.h"
-
-#include <csignal>
+#include "plugins/interpolate_plugin.h"
+#include "NvInfer.h"
+#include "NvInferRuntimeCommon.h"
 
 namespace trtorch {
 namespace core {
@@ -11,6 +12,54 @@ namespace converters {
 namespace impl {
 namespace {
 
+/*
+ * Helper functions
+ */
+
+void create_plugin(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in, const char* name, 
+                                                                                         std::vector<int64_t> in_shape, 
+                                                                                         std::vector<int64_t> out_shape, 
+                                                                                         std::vector<int64_t> out_size, 
+                                                                                         std::string mode) {
+    LOG_WARNING("Interpolation layer will be run through ATen, not TensorRT. Performance may differ.");
+    
+    auto creator = new plugins::InterpolatePluginCreator();
+    auto plugin = creator->createPlugin(name, in_shape, out_shape, out_size, mode, false);
+
+    auto resize_layer = ctx->net->addPluginV2(reinterpret_cast<nvinfer1::ITensor* const*>(&in), 1, *plugin);
+    TRTORCH_CHECK(resize_layer, "Unable to create interpolation plugin from node" << *n);
+
+    resize_layer->setName(util::node_info(n).c_str());
+
+    auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
+
+    LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
+}
+
+void resize_layer_size(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* in, std::vector<int64_t> out_shape, 
+                                                                                             nvinfer1::ResizeMode mode) {
+    auto resize_layer = ctx->net->addResize(*in);
+    TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
+
+    resize_layer->setOutputDimensions(util::toDims(out_shape));
+    resize_layer->setResizeMode(mode);
+    resize_layer->setName(util::node_info(n).c_str());
+    
+    // if interpolation mode is linear, align corners must have been set to true. else, don't use align corners.
+    if (mode == nvinfer1::ResizeMode::kLINEAR) {
+        resize_layer->setAlignCorners(true);
+    }
+
+    auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
+    
+    LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
+}
+
+
+/*
+ * Interpolate Converter
+ */
+
 auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
     .pattern({
         "aten::upsample_nearest1d(Tensor self, int[1] output_size, float? scales=None) -> (Tensor)",
@@ -27,15 +76,7 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
-                auto resize_layer = ctx->net->addResize(*in);
-                TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
-
-                resize_layer->setOutputDimensions(util::toDims(out_shape));
-                resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
-                resize_layer->setName(util::node_info(n).c_str());
-
-                auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
-                LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
+                resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kNEAREST);
             } else {
                 TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest1d not supported yet.");
             }
@@ -57,15 +98,7 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
-                auto resize_layer = ctx->net->addResize(*in);
-                TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
-
-                resize_layer->setOutputDimensions(util::toDims(out_shape));
-                resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
-                resize_layer->setName(util::node_info(n).c_str());
-
-                auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
-                LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
+                resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kNEAREST);
             } else {
                 TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest2d not supported yet.");
             }
@@ -86,18 +119,94 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
 
                 auto out_shape = in_shape;
                 std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
+
+                resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kNEAREST);
+            } else {
+                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest3d not supported yet.");
+            }
+
+            return true;
+        }
+    }).pattern({
+        "aten::upsample_linear1d(Tensor self, int[1] output_size, bool align_corners, float? scales=None) -> (Tensor)",
+        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+            auto in = args[0].ITensor();
+            auto in_shape = util::toVec(in->getDimensions());
+            bool align_corners = args[2].unwrapToBool();
+
+            // Case 1: user uses output size and not scales
+            if (!args[1].IValue()->isNone() && args[3].IValue()->isNone()) {
+                auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
+
+                TRTORCH_ASSERT(out_size.size() == 1, "aten::upsample_linear1d input Tensor and output size dimension mismatch");
 
-                auto resize_layer = ctx->net->addResize(*in);
-                TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
+                auto out_shape = in_shape; 
+                std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
 
-                resize_layer->setOutputDimensions(util::toDims(out_shape));
-                resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
-                resize_layer->setName(util::node_info(n).c_str());
+                if (!align_corners) {
+                    // align_corners not supported in TensorRT, create plugin and run layer through PyTorch
+                    create_plugin(ctx, n, in, "linear1d", in_shape, out_shape, out_size, std::string("linear"));
+                } else {
+                    resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR);
+                }
+            } else {
+                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_linear1d not supported yet.");
+            }
+
+            return true;
+        }
+    }).pattern({
+        "aten::upsample_bilinear2d(Tensor self, int[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> (Tensor)",
+        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+            auto in = args[0].ITensor();
+            auto in_shape = util::toVec(in->getDimensions());
+            bool align_corners = args[2].unwrapToBool();
 
-                auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], resize_layer->getOutput(0));
-                LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
+            // Case 1: user uses output size and not scales_h, scales_w
+            if (!args[1].IValue()->isNone() && args[3].IValue()->isNone() && args[4].IValue()->isNone()) {
+                auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
+
+                TRTORCH_ASSERT(out_size.size() == 2, "aten::upsample_bilinear2d input Tensor and output size dimension mismatch");
+                
+                auto out_shape = in_shape;
+                std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
+
+                if (!align_corners) {
+                    // align_corners not supported in TensorRT, create plugin and run layer through PyTorch
+                    create_plugin(ctx, n, in, "bilinear2d", in_shape, out_shape, out_size, std::string("bilinear"));
+                } else {
+                    resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR);
+                }
             } else {
-                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest3d not supported yet.");
+                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_bilinear2d not supported yet.");
+            }
+
+            return true;
+        }
+    }).pattern({
+        "aten::upsample_trilinear3d(Tensor self, int[3] output_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> (Tensor)",
+        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+            auto in = args[0].ITensor();
+            auto in_shape = util::toVec(in->getDimensions());
+            bool align_corners = args[2].unwrapToBool();
+
+            // Case 1: user uses output size and not scales_d, scales_h, scales_w
+            if (!args[1].IValue()->isNone() && args[3].IValue()->isNone() && args[4].IValue()->isNone() && args[5].IValue()->isNone()) {
+                auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
+
+                TRTORCH_ASSERT(out_size.size() == 3, "aten::upsample_trilinear3d input Tensor and output size dimension mismatch");
+                
+                auto out_shape = in_shape;
+                std::copy(out_size.begin(), out_size.end(), out_shape.begin() + (in_shape.size() - out_size.size()));
+
+                if (!align_corners) {
+                    // align_corners not supported in TensorRT, create plugin and run layer through PyTorch
+                    create_plugin(ctx, n, in, "trilinear3d", in_shape, out_shape, out_size, std::string("trilinear"));
+                } else {
+                    resize_layer_size(ctx, n, in, out_shape, nvinfer1::ResizeMode::kLINEAR);
+                }
+            } else {
+                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_trilinear3d not supported yet.");
             }
 
             return true;
@@ -110,4 +219,4 @@ auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // namespace trtorch
+} // namespace trtorch

core/conversion/converters/impl/plugins/BUILD

@@ -0,0 +1,35 @@
+package(default_visibility = ["//visibility:public"])
+
+config_setting(
+    name = "use_pre_cxx11_abi",
+    values = {
+        "define": "abi=pre_cxx11_abi",
+    }
+)
+
+cc_library(
+    name = "plugins",
+    hdrs = [
+        "interpolate_plugin.h"
+    ],
+    srcs = [
+        "interpolate_plugin.cpp"
+    ],
+    deps = [
+        "@tensorrt//:nvinfer",
+        "//core/util:prelude",
+        "//core/conversion/conversionctx",
+    ] + select({
+        ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
+        "//conditions:default":  ["@libtorch//:libtorch"],
+    }),
+    alwayslink = True,
+)
+
+load("@rules_pkg//:pkg.bzl", "pkg_tar")
+
+pkg_tar(
+    name = "include",
+    package_dir = "core/conversion/converters/impl/plugins",
+    srcs = ["interpolate_plugin.h"],
+)