Merge pull request #362 from NVIDIA/bowa_core_elu

add ELU converter to core library

Author: narendasan

core/conversion/converters/impl/activation.cpp

@@ -152,6 +152,21 @@ auto acthardtanh TRTORCH_UNUSED =
                     out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
                     LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
                     return true;
+                  }})
+        .pattern({"aten::elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    auto in = args[0].ITensorOrFreeze(ctx);
+                    auto alpha = args[1].unwrapToDouble();
+
+                    auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kELU);
+                    TRTORCH_CHECK(new_layer, "Unable to create layer for aten::elu");
+                    new_layer->setAlpha(alpha);
+
+                    new_layer->setName(trtorch::core::util::node_info(n).c_str());
+
+                    auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
+                    LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
+                    return true;
                   }});
 
 } // namespace

tests/core/conversion/converters/test_activation.cpp

@@ -175,5 +175,28 @@ TEST(Converters, ATenLeakyReluConvertsCorrectly) {
   params = trtorch::core::conversion::get_named_params(g->inputs(), {});
   auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
 
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
+}
+
+TEST(Converters, ATenEluConvertsCorrectly) {
+  const auto graph = R"IR(
+       graph(%x.1 : Tensor):
+        %2 : float = prim::Constant[value=1.]()
+        %3 : int = prim::Constant[value=1]()
+        %result.2 : Tensor = aten::elu(%x.1, %2, %3, %3)
+        return (%result.2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(-5, 5, {1, 10, 1, 1}, {at::kCUDA});
+
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  in = at::clone(in);
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0], 2e-6));
 }