feat(//core/conversion/converters/impl): logic implemented

Signed-off-by: Abhiram Iyer <abhi.iyer.ai@gmail.com> Signed-off-by: Abhiram Iyer <abhirami@nvidia.com>
pytorch · Jun 8, 2020 · 7f12160 · 7f12160
1 parent 1b50484
commit 7f12160
Showing 1 changed file with 16 additions and 6 deletions.
diff --git a/core/conversion/converters/impl/interpolate.cpp b/core/conversion/converters/impl/interpolate.cpp
@@ -25,6 +25,12 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
                 auto output_size = util::toDims(args[1].unwrapToIntList());
 
                 TRTORCH_ASSERT(output_size.nbDims == 1, "aten::upsample_nearest1d 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);
+
+                resize_layer->setOutputDimensions(output_size);
+                resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
             } else {
                 LOG_DEBUG("scale factor parameters not supported yet.");
             }
@@ -34,7 +40,8 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
     }).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 {
-//            std::raise(SIGINT);
+            // std::raise(SIGINT);
+
             TRTORCH_ASSERT(args[0].IValue()->isTensor(), "Input expected to be of type Tensor");
 
             auto in = args[0].ITensor();
@@ -46,12 +53,11 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
 
                 TRTORCH_ASSERT( (output_size.nbDims == 1 || output_size.nbDims == 2), "aten::upsample_nearest2d input Tensor and output size dimension mismatch");
 
-                nvinfer1::ILayer* new_layer;
-
-
-
-                //util::toDims(args[1].unwrapToIntList());
+                auto resize_layer = ctx->net->addResize(*in);
+                TRTORCH_CHECK(resize_layer, "Unable to create interpolation (resizing) layer from node" << *n);
 
+                resize_layer->setOutputDimensions(output_size);
+                resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
             } else {
                 LOG_DEBUG("scale factor parameters not supported yet.");
             }
@@ -72,7 +78,11 @@ auto interpolate_registrations = RegisterNodeConversionPatterns()
 
                 TRTORCH_ASSERT( (output_size.nbDims == 1 || output_size.nbDims == 3), "aten::upsample_nearest3d 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);
 
+                resize_layer->setOutputDimensions(output_size);
+                resize_layer->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
             } else {
                 LOG_DEBUG("scale factor parameters not supported yet.");
             }