Merge pull request #374 from NVIDIA/yutec_aten_tranaspose

Add aten.transpose support.

Author: narendasan

core/conversion/converters/impl/shuffle.cpp

@@ -92,8 +92,40 @@ static auto shuffle_registrations TRTORCH_UNUSED =
                     auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
                     LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
 
+                    return true;
+                  }})
+        .pattern({"aten::transpose.int(Tensor(a) self, int dim0, int dim1) -> (Tensor(a))",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto in = args[0].ITensorOrFreeze(ctx);
+                    auto in_shape = util::toVec(in->getDimensions());
+                    auto ndims = in_shape.size();
+                    auto dim0 = args[1].unwrapToInt();
+                    auto dim1 = args[2].unwrapToInt();
+
+                    std::vector<int64_t> new_order;
+                    for (size_t i = 0; i < ndims; i++) {
+                      new_order.push_back(i);
+                    }
+                    auto tmp = dim0;
+                    new_order[dim0] = new_order[dim1];
+                    new_order[dim1] = tmp;
+
+                    LOG_DEBUG("Shuffle to: " << util::toDims(new_order));
+
+                    auto shuffle = ctx->net->addShuffle(*in);
+                    TRTORCH_CHECK(shuffle, "Unable to create shuffle layer from node: " << *n);
+                    nvinfer1::Permutation permute;
+                    std::copy(new_order.begin(), new_order.end(), permute.order);
+
+                    shuffle->setSecondTranspose(permute);
+                    shuffle->setName(util::node_info(n).c_str());
+
+                    auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
+                    LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
+
                     return true;
                   }});
+
 } // namespace
 } // namespace impl
 } // namespace converters

tests/core/conversion/converters/test_shuffle.cpp

@@ -214,3 +214,29 @@ TEST(Converters, ATenFlattenConvertsCorrectlyWithDynamicBatch) {
 
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
 }
+
+TEST(Converters, ATenTransposeConvertsCorrectly) {
+  const auto graph = R"IR(
+    graph(%x.1 : Tensor):
+      %2 : int = prim::Constant[value=1]()
+      %3 : int = prim::Constant[value=3]()
+      %4 : Tensor = aten::transpose(%x.1, %2, %3)
+      return (%4))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(0, 5, {2, 3, 4, 5, 6}, {at::kCUDA});
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+
+  std::cout << "Running JIT" << std::endl;
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  std::cout << "Running TRT" << std::endl;
+  in = at::clone(in);
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+  auto trt = trt_results[0].reshape_as(jit_results[0]);
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt, 2e-6));
+}