Add support for sigmoid and tanh

Signed-off-by: Junjie Bai <baidingding7@gmail.com>

Author: bddppq

core/conversion/converters/impl/activation.cpp

@@ -9,37 +9,42 @@ namespace converters {
 namespace impl {
 namespace {
 
-bool relu(ConversionCtx* ctx, const torch::jit::Node* n, args& args) {
-     auto in = args[0].ITensor();
+#define convert(act, trt_type)                                                 \
+  bool act(ConversionCtx* ctx, const torch::jit::Node* n, args& args) {        \
+    auto in = args[0].ITensor();                                               \
+                                                                               \
+    auto new_layer =                                                           \
+        ctx->net->addActivation(*in, nvinfer1::ActivationType::trt_type);      \
+    TRTORCH_CHECK(new_layer,                                                   \
+                  "Unable to create " #act " layer from node: " << *n);        \
+                                                                               \
+    new_layer->setName(util::node_info(n).c_str());                            \
+    auto out_value = n->outputs()[0];                                          \
+    auto out_tensor = new_layer->getOutput(0);                                 \
+    out_tensor->setName(out_value->debugName().c_str());                       \
+    ctx->value_tensor_map[out_value] = out_tensor;                             \
+    LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());         \
+                                                                               \
+    return true;                                                               \
+  }                                                                            \
+                                                                               \
+  auto act##_registrations TRTORCH_UNUSED =                                    \
+      RegisterNodeConversionPatterns()                                         \
+          .pattern({"aten::" #act "(Tensor input) -> (Tensor)",                \
+                    [](ConversionCtx *ctx, const torch::jit::Node *n,          \
+                       args &args) -> bool { return act(ctx, n, args); }})     \
+          .pattern({"aten::" #act "_(Tensor(a!) self) -> (Tensor(a!))",        \
+                    [](ConversionCtx *ctx, const torch::jit::Node *n,          \
+                       args &args) -> bool { return act(ctx, n, args); }});
 
-     auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kRELU);
-     TRTORCH_CHECK(new_layer, "Unable to create ReLU layer from node: " << *n);
+convert(relu, kRELU);
+convert(sigmoid, kSIGMOID);
+convert(tanh, kTANH);
 
-     new_layer->setName(util::node_info(n).c_str());
-     auto out_value = n->outputs()[0];
-     auto out_tensor = new_layer->getOutput(0);
-     out_tensor->setName(out_value->debugName().c_str());
-     ctx->value_tensor_map[out_value] = out_tensor;
-     LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
-
-     return true;
-}
-
-auto relu_registrations = RegisterNodeConversionPatterns()
-    .pattern({
-        "aten::relu(Tensor input) -> (Tensor)",
-        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-            return relu(ctx, n, args);
-        }
-    }).pattern({
-        "aten::relu_(Tensor(a!) self) -> (Tensor(a!))",
-        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
-            return relu(ctx, n, args);
-        }
-    });
+#undef convert
 } // namespace
 } // namespace impl
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // trtorch 
+} // namespace trtorch

core/util/macros.h

@@ -53,3 +53,11 @@
       TRTORCH_THROW_ERROR("Expected " << #cond                              \
                           << " to be true but got false\n" << __VA_ARGS__); \
   }
+
+
+// suppress an unused variable.
+#if defined(_MSC_VER) && !defined(__clang__)
+#define TRTORCH_UNUSED __pragma(warning(suppress: 4100 4101))
+#else
+#define TRTORCH_UNUSED __attribute__((__unused__))
+#endif //_MSC_VER

tests/core/converters/BUILD

@@ -5,7 +5,7 @@ converter_test(
 )
 
 converter_test(
-    name = "test_relu"
+    name = "test_activation"
 )
 
 converter_test(
@@ -36,7 +36,7 @@ test_suite(
     name = "test_converters",
     tests = [
         ":test_softmax",
-        ":test_relu",
+        ":test_activation",
         ":test_pooling",
         ":test_unary",
         ":test_linear",

tests/core/converters/test_activation.cpp

@@ -0,0 +1,65 @@
+#include <string>
+#include "gtest/gtest.h"
+#include "torch/csrc/jit/irparser.h"
+#include "tests/util/util.h"
+#include "core/compiler.h"
+
+TEST(Converters, ATenReLUConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %3 : Tensor = aten::relu(%0)
+        return (%3))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::script::parseIR(graph, &*g);
+
+    auto in = at::randint(-5, 5, {5}, {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]));
+}
+
+TEST(Converters, ATenSigmoidConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %3 : Tensor = aten::sigmoid(%0)
+        return (%3))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::script::parseIR(graph, &*g);
+
+    auto in = at::randint(-5, 5, {5}, {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]));
+}
+
+TEST(Converters, ATenTanhConvertsCorrectly) {
+    const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %3 : Tensor = aten::tanh(%0)
+        return (%3))IR";
+
+    auto g = std::make_shared<torch::jit::Graph>();
+    torch::jit::script::parseIR(graph, &*g);
+
+    auto in = at::randint(-5, 5, {5}, {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]));
+}