fix bugs in split converter

Signed-off-by: Ruoqian Guo <[email protected]>

Author: ruoqianguo

core/conversion/converters/impl/select.cpp

@@ -18,7 +18,7 @@ bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool s
   auto in = args[0].ITensor();
   auto axis = args[2].unwrapToInt();
   auto inDimSize = in->getDimensions().d[axis];
-  auto numOutputs = 1;
+  auto numOutputs = 1, numRemainder = 0;
   std::vector<int64_t> sizes;
 
   if (split_list) {
@@ -27,10 +27,13 @@ bool add_split(ConversionCtx* ctx, const torch::jit::Node* n, args& args, bool s
   } else {
     auto split_size = args[1].unwrapToInt();
     numOutputs = inDimSize / split_size;
-    if (numOutputs == 1) {
+    numRemainder = inDimSize % split_size;
+    for (int i = 0; i < numOutputs; i++) {
       sizes.push_back(split_size);
-    } else {
-      sizes = std::vector<int64_t>(numOutputs, 1);
+    }
+    if (numRemainder) {
+      numOutputs += 1;
+      sizes.push_back(numRemainder);
     }
   }
 

core/conversion/var/BUILD

@@ -19,7 +19,8 @@ cc_library(
     deps = [
         "@tensorrt//:nvinfer",
         "//core/util:prelude",
-        "//core/conversion/converters:weights"
+        "//core/conversion/converters:weights",
+        "//core/conversion/tensorcontainer:tensorcontainer"
     ] + select({
         ":use_pre_cxx11_abi":  ["@libtorch_pre_cxx11_abi//:libtorch"],
         "//conditions:default":  ["@libtorch//:libtorch"],

core/conversion/var/Var.cpp

@@ -90,24 +90,37 @@ nvinfer1::ITensor* Var::ITensorOrFreeze(ConversionCtx* ctx) {
     LOG_DEBUG(ctx->logger, "Found IValue containing object of type " << *(ptr_.ivalue->type()));
   }
   TRTORCH_CHECK(
-      isITensor() || (isIValue() && ptr_.ivalue->isTensor()),
+      isITensor() || (isIValue() && (ptr_.ivalue->isTensor() || ptr_.ivalue->isCustomClass())),
       "Requested either IValue containing a Tensor, or ITensor, however Var type is " << type_name());
 
   nvinfer1::ITensor* out;
-
+  auto weights = converters::Weights();
   if (isIValue()) {
-    auto weights = converters::Weights(ctx, ptr_.ivalue->toTensor());
-
-    auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
-    TRTORCH_CHECK(const_layer, "Unable to freeze tensor into constant layer");
-
-    out = const_layer->getOutput(0);
-
-    std::ostringstream tensor_id;
-    tensor_id << reinterpret_cast<int*>(out);
-
-    LOG_DEBUG(ctx->logger, "Freezing tensor " << tensor_id.str() << " as an IConstantLayer");
-    const_layer->setName(("[Freeze Tensor " + tensor_id.str() + " ]").c_str());
+    if (ptr_.ivalue->isTensor()) {
+      auto tensor = ptr_.ivalue->toTensor();
+      if (tensor.scalar_type() == at::kLong) {
+        weights = converters::Weights(ctx, tensor.toType(at::kInt));
+      } else if (tensor.scalar_type() == at::kDouble) {
+        weights = converters::Weights(ctx, tensor.toType(at::kFloat));
+      } else {
+        weights = converters::Weights(ctx, tensor);
+      }
+
+      auto const_layer = ctx->net->addConstant(weights.shape, weights.data);
+      TRTORCH_CHECK(const_layer, "Unable to freeze tensor into constant layer");
+
+      out = const_layer->getOutput(0);
+
+      std::ostringstream tensor_id;
+      tensor_id << reinterpret_cast<int*>(out);
+
+      LOG_DEBUG(ctx->logger, "Freezing tensor " << tensor_id.str() << " as an IConstantLayer");
+      const_layer->setName(("[Freeze Tensor " + tensor_id.str() + " ]").c_str());
+    } else {
+      // Split converter generates c10::IValue which hold TensorContainer.
+      auto output_container = ptr_.ivalue->toCustomClass<TensorContainer>();
+      out = output_container.get()->tensor();
+    }
   } else {
     out = ptr_.tensor;
   }

core/conversion/var/Var.h

@@ -5,6 +5,7 @@
 
 #include "core/conversion/conversionctx/ConversionCtx.h"
 #include "core/conversion/converters/Weights.h"
+#include "core/conversion/tensorcontainer/TensorContainer.h"
 #include "core/util/prelude.h"
 #include "torch/csrc/jit/ir/ir.h"
 

tests/core/conversion/converters/test_select.cpp

@@ -288,3 +288,60 @@ TEST(Converters, ATenSplitFixedConvertsCorrectly) {
     ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[i], trt, 2e-6));
   }
 }
+
+TEST(Converters, ATenSplitFixedHasRemainderConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%argument_1.1 : Tensor):
+          %2 : int = prim::Constant[value=2]()
+          %2.1 : int = prim::Constant[value=1]()
+          %3 : Tensor[] = aten::split(%argument_1.1, %2, %2.1)
+          %4 : Tensor, %5 : Tensor, %6 : Tensor = prim::ListUnpack(%3)
+          return (%4, %5, %6))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {1, 5, 4, 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);
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+  for (size_t i = 0; i < jit_results.size(); i++) {
+    auto trt = trt_results[i].reshape(jit_results[i].sizes());
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[i], trt, 2e-6));
+  }
+}
+
+TEST(Converters, ATenSplitAndAddConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%argument_1.1 : Tensor):
+          %2 : int = prim::Constant[value=2]()
+          %2.1 : int = prim::Constant[value=1]()
+          %3 : Tensor[] = aten::split(%argument_1.1, %2, %2.1)
+          %4 : Tensor, %5 : Tensor = prim::ListUnpack(%3)
+          %6 : Tensor = aten::add(%4, %5, %2.1)
+          return (%6))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {1, 4, 4, 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);
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {trt_in});
+
+  for (size_t i = 0; i < jit_results.size(); i++) {
+    auto trt = trt_results[i].reshape(jit_results[i].sizes());
+    ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[i], trt, 2e-6));
+  }
+}