Support for interpolation (aten::upsample_nearest)

core/conversion/InterfaceTypes.cpp

@@ -55,9 +55,9 @@ InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_
                   << max_shape.size() << ")");
     }
 
-    min = util::toDimsPad(min_shape, 4);
-    opt = util::toDimsPad(opt_shape, 4);
-    max = util::toDimsPad(max_shape, 4);
+    min = util::toDims(min_shape);
+    opt = util::toDims(opt_shape);
+    max = util::toDims(max_shape);
 
     std::vector<int64_t> dyn_shape;
     for (size_t i = 0; i < opt_shape.size(); i++) {
@@ -72,7 +72,7 @@ InputRange::InputRange(std::vector<int64_t> min_shape, std::vector<int64_t> opt_
         }
     }
 
-    input_shape = util::toDimsPad(dyn_shape, 4);
+    input_shape = util::toDims(dyn_shape);
 
 }
 

core/conversion/converters/BUILD

@@ -29,6 +29,7 @@ cc_library(
         "impl/shuffle.cpp",
         "impl/softmax.cpp",
         "impl/unary.cpp",
+        "impl/interpolate.cpp"
     ],
     deps = [
         "@tensorrt//:nvinfer",

core/conversion/converters/impl/interpolate.cpp

@@ -0,0 +1,113 @@
+#include "torch/torch.h"
+#include "core/util/prelude.h"
+#include "core/conversion/converters/converters.h"
+
+#include <csignal>
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto interpolate_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
+    .pattern({
+        "aten::upsample_nearest1d(Tensor self, int[1] output_size, 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());
+
+            // Case 1: user uses output size and not scales
+            if (!args[1].IValue()->isNone() && args[2].IValue()->isNone()) {
+                auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
+
+                TRTORCH_ASSERT(out_size.size() == 1, "aten::upsample_nearest1d 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()));
+
+                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());
+            } else {
+                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest1d not supported yet.");
+            }
+
+            return true;
+        }
+    }).pattern({
+        "aten::upsample_nearest2d(Tensor self, int[2] output_size, 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());
+
+            // Case 1: user uses output_size and not scales_h, scales_w
+            if (!args[1].IValue()->isNone() && args[2].IValue()->isNone() && args[3].IValue()->isNone()){
+                auto out_size = util::toVec(util::toDims(args[1].unwrapToIntList()));
+
+                TRTORCH_ASSERT(out_size.size() == 2, "aten::upsample_nearest2d 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()));
+
+                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());
+            } else {
+                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest2d not supported yet.");
+            }
+
+            return true;
+        }
+    }).pattern({
+        "aten::upsample_nearest3d(Tensor self, int[3] output_size, 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());
+
+            // Case 1: user uses output size and not scales_d, scales_h, scales_w
+            if (!args[1].IValue()->isNone() && args[2].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() == 3, "aten::upsample_nearest3d 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()));
+
+                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());
+            } else {
+                TRTORCH_THROW_ERROR("Unable to convert node: " << util::node_info(n) << "\nScale factor parameter for upsample_nearest3d not supported yet.");
+            }
+
+            return true;
+        }
+    });
+
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch

tests/core/converters/BUILD

@@ -51,6 +51,10 @@ converter_test(
   name = "test_unary"
 )
 
+converter_test(
+  name = "test_interpolate"
+)
+
 test_suite(
   name = "test_converters",
   tests = [
@@ -65,6 +69,7 @@ test_suite(
     ":test_shuffle",
     ":test_softmax",
     ":test_unary",
+    ":test_interpolate",
   ]
 )
 

tests/core/converters/test_interpolate.cpp

@@ -0,0 +1,150 @@
+#include <string>
+#include "gtest/gtest.h"
+#include "torch/csrc/jit/ir/irparser.h"
+#include "tests/util/util.h"
+#include "core/compiler.h"
+
+TEST(Converters, ATenUpsampleNearest1dConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=10]()
+        %3 : int[] = prim::ListConstruct(%2)
+        %4 : None = prim::Constant()
+        %5 : Tensor = aten::upsample_nearest1d(%0, %3, %4)
+        return (%5))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    // Input Tensor needs to be 3D for TensorRT upsample_nearest1d
+    auto in = at::randint(1, 10, {10, 2, 2}, {at::kCUDA});
+
+    auto jit_in = at::clone(in);
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+    auto trt_in = at::clone(in);
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+    auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenUpsampleNearest2dConvertsCorrectly1dOutputSize) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=10]()
+        %3 : int[] = prim::ListConstruct(%2, %2)
+        %4 : None = prim::Constant()
+        %5 : Tensor = aten::upsample_nearest2d(%0, %3, %4, %4)
+        return (%5))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    // Input Tensor needs to be 4D for TensorRT upsample_nearest2d
+    auto in = at::randint(1, 10, {10, 2, 2, 2}, {at::kCUDA});
+
+    auto jit_in = at::clone(in);
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+    auto trt_in = at::clone(in);
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+    auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenUpsampleNearest2dConvertsCorrectly2dOutputSize) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=10]()
+        %3 : int[] = prim::ListConstruct(%2, %2)
+        %4 : None = prim::Constant()
+        %5 : Tensor = aten::upsample_nearest2d(%0, %3, %4, %4)
+        return (%5))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    // Input Tensor needs to be 4D for TensorRT upsample_nearest2d
+    auto in = at::randint(1, 10, {10, 2, 2, 2}, {at::kCUDA});
+
+    auto jit_in = at::clone(in);
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+    auto trt_in = at::clone(in);
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+    auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenUpsampleNearest3dConvertsCorrectly1dOutputSize) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=10]()
+        %3 : int[] = prim::ListConstruct(%2, %2, %2)
+        %4 : None = prim::Constant()
+        %5 : Tensor = aten::upsample_nearest3d(%0, %3, %4, %4, %4)
+        return (%5))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    // Input Tensor needs to be 5D for TensorRT upsample_nearest3d
+    auto in = at::randint(1, 10, {10, 2, 2, 2, 2}, {at::kCUDA});
+
+    auto jit_in = at::clone(in);
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+    auto trt_in = at::clone(in);
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+    auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenUpsampleNearest3dConvertsCorrectly3dOutputSize) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %2 : int = prim::Constant[value=10]()
+        %3 : int[] = prim::ListConstruct(%2, %2, %2)
+        %4 : None = prim::Constant()
+        %5 : Tensor = aten::upsample_nearest3d(%0, %3, %4, %4, %4)
+        return (%5))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    // Input Tensor needs to be 5D for TensorRT upsample_nearest3d
+    auto in = at::randint(1, 10, {10, 2, 2, 2, 2}, {at::kCUDA});
+
+    auto jit_in = at::clone(in);
+    auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto jit_results = trtorch::tests::util::RunGraph(g, params, {jit_in});
+
+    auto trt_in = at::clone(in);
+    params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+    auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+    auto trt = trt_results[0].reshape(jit_results[0].sizes());
+
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}