Merge pull request #188 from peri044/power_layer

feat(//core/converters): Add conversion support for torch.narrow()

Author: narendasan

core/conversion/converters/impl/select.cpp

@@ -46,6 +46,77 @@ auto select_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns()
 
             LOG_DEBUG("Output tensor shape: " << out->getDimensions());
 
+            return true;
+        }
+    }).pattern({
+        "aten::narrow(Tensor(a) self, int dim, int start, int length) -> Tensor(a)",
+        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+            auto in = args[0].ITensor();
+            auto axis  = args[1].unwrapToInt();
+            auto start = (int32_t) args[2].unwrapToInt();
+            auto length = (int32_t) args[3].unwrapToInt();
+
+            // index to access needs to be an at::Tensor
+            at::Tensor indices = torch::arange(start, start + length, 1).to(torch::kI32);
+            auto weights = Weights(ctx, indices);
+            
+            // IConstantLayer to convert indices from Weights to ITensor
+            auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
+            TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+            auto const_out = const_layer->getOutput(0);
+            
+            // IGatherLayer takes in input tensor, the indices, and the axis of input tensor to take indices from
+            auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
+            TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
+            auto gather_out = gather_layer->getOutput(0);
+            
+            // IShuffleLayer removes redundant dimensions
+            auto shuffle_layer = ctx->net->addShuffle(*gather_out);
+            TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+            shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+            shuffle_layer->setName(util::node_info(n).c_str());
+            auto shuffle_out = shuffle_layer->getOutput(0);
+
+            auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
+
+            LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+            return true;
+        }
+    }).pattern({
+        "aten::narrow.Tensor(Tensor(a) self, int dim, Tensor start, int length) -> Tensor(a)",
+        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+            auto in = args[0].ITensor();
+            auto axis  = args[1].unwrapToInt();
+            torch::Tensor start = args[2].IValue()->toTensor().to(torch::kI32);
+            int32_t startIdx = start.item().to<int32_t>();
+            auto length = (int32_t) args[3].unwrapToInt();
+
+            // index to access needs to be an at::Tensor
+            at::Tensor indices = torch::arange(startIdx, startIdx + length, 1).to(torch::kI32);
+            auto weights = Weights(ctx, indices);
+            
+            // IConstantLayer to convert indices from Weights to ITensor
+            auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
+            TRTORCH_CHECK(const_layer, "Unable to create constant layer from node: " << *n);
+            auto const_out = const_layer->getOutput(0);
+            
+            // IGatherLayer takes in input tensor, the indices, and the axis of input tensor to take indices from
+            auto gather_layer = ctx->net->addGather(*in, *const_out, axis);
+            TRTORCH_CHECK(gather_layer, "Unable to create gather layer from node: " << *n);
+            auto gather_out = gather_layer->getOutput(0);
+            
+            // IShuffleLayer removes redundant dimensions
+            auto shuffle_layer = ctx->net->addShuffle(*gather_out);
+            TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+            shuffle_layer->setReshapeDimensions(util::unpadDims(gather_out->getDimensions()));
+            shuffle_layer->setName(util::node_info(n).c_str());
+            auto shuffle_out = shuffle_layer->getOutput(0);
+
+            auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_out);
+
+            LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
             return true;
         }
     });

tests/core/converters/test_select.cpp

@@ -55,5 +55,32 @@ TEST(Converters, ATenSelectIntTwiceConvertsCorrectly) {
 
     auto trt = trt_results[0].reshape(jit_results[0].sizes());
 
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}
+
+TEST(Converters, ATenNarrowStartScalarConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%x.1 : Tensor):
+            %2 : int = prim::Constant[value=2]()
+            %3 : int = prim::Constant[value=0]()
+            %4 : Tensor = aten::narrow(%x.1, %3, %3, %2)
+            return (%4))IR";
+    
+    auto g = std::make_shared<torch::jit::Graph>();
+
+    torch::jit::parseIR(graph, &*g);
+
+    auto in = at::randint(1, 10, {3, 2, 2, 4}, {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));
 }