Merge pull request #329 from inocsin/clamp_max_min

add clamp_min/clamp_max converter

Author: narendasan

core/conversion/converters/impl/element_wise.cpp

@@ -68,6 +68,21 @@ nvinfer1::ILayer* add_elementwise(
   return ele;
 }
 
+nvinfer1::ITensor* clamp_util(
+    ConversionCtx* ctx,
+    const torch::jit::Node* n,
+    nvinfer1::ITensor* self,
+    float limit,
+    nvinfer1::ElementWiseOperation op_type,
+    std::string str) {
+  nvinfer1::ITensor* clamp_layer_out = self;
+  auto limitTensor = tensor_to_const(ctx, torch::tensor({limit}));
+  auto limit_layer = add_elementwise(ctx, op_type, clamp_layer_out, limitTensor, util::node_info(n) + str);
+  TRTORCH_CHECK(limit_layer, "Unable to create elementwise " << str << " layer for node: " << *n);
+  clamp_layer_out = limit_layer->getOutput(0);
+  return clamp_layer_out;
+}
+
 auto element_wise_registrations TRTORCH_UNUSED =
     RegisterNodeConversionPatterns()
         .pattern({"aten::add.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> "
@@ -145,38 +160,58 @@ auto element_wise_registrations TRTORCH_UNUSED =
                     return true;
                   }})
         .pattern({"aten::clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    // Compute min(max(min_threshold, input), max_threshold)
+                    auto self = args[0].ITensorOrFreeze(ctx);
+                    auto clamp_layer_out = self;
+
+                    if (args[1].isIValue() && args[1].IValue()->isScalar() && args[2].isIValue() &&
+                        args[2].IValue()->isScalar()) {
+                      auto alpha = args[1].unwrapToScalar().to<float>();
+                      auto beta = args[2].unwrapToScalar().to<float>();
+                      auto clip_layer = ctx->net->addActivation(*self, nvinfer1::ActivationType::kCLIP);
+                      TRTORCH_CHECK(clip_layer, "Unable to create clip layer for node: " << *n);
+                      clip_layer->setAlpha(alpha);
+                      clip_layer->setBeta(beta);
+                      clamp_layer_out = clip_layer->getOutput(0);
+                    } else if (args[1].isIValue() && args[1].IValue()->isScalar()) {
+                      auto limit = args[1].unwrapToScalar().to<float>();
+                      clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max");
+                    } else if (args[2].isIValue() && args[2].IValue()->isScalar()) {
+                      auto limit = args[2].unwrapToScalar().to<float>();
+                      clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min");
+                    }
+
+                    auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
+                    LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions());
+                    return true;
+                  }})
+        .pattern({"aten::clamp_min(Tensor self, Scalar min) -> (Tensor)",
                   [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                     // Compute min(max(min_threshold, input), max_threshold)
                     auto self = args[0].ITensorOrFreeze(ctx);
                     auto clamp_layer_out = self;
                     if (args[1].isIValue() && args[1].IValue()->isScalar()) {
-                      auto minScalar = args[1].unwrapToScalar().to<float>();
-                      auto minTensor = tensor_to_const(ctx, torch::tensor({minScalar}));
-                      auto max_layer = add_elementwise(
-                          ctx,
-                          nvinfer1::ElementWiseOperation::kMAX,
-                          clamp_layer_out,
-                          minTensor,
-                          util::node_info(n) + std::string("_max"));
-                      TRTORCH_CHECK(max_layer, "Unable to create elementwise max layer for node: " << *n);
-                      clamp_layer_out = max_layer->getOutput(0);
+                      auto limit = args[1].unwrapToScalar().to<float>();
+                      clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMAX, "_max");
                     }
 
-                    if (args[2].isIValue() && args[2].IValue()->isScalar()) {
-                      auto maxScalar = args[2].unwrapToScalar().to<float>();
-                      auto maxTensor = tensor_to_const(ctx, torch::tensor({maxScalar}));
-                      auto min_layer = add_elementwise(
-                          ctx,
-                          nvinfer1::ElementWiseOperation::kMIN,
-                          clamp_layer_out,
-                          maxTensor,
-                          util::node_info(n) + std::string("_min"));
-                      TRTORCH_CHECK(min_layer, "Unable to create elementwise min layer for node: " << *n);
-                      clamp_layer_out = min_layer->getOutput(0);
+                    auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
+                    LOG_DEBUG("clamp_min layer output tensor shape: " << out->getDimensions());
+                    return true;
+                  }})
+        .pattern({"aten::clamp_max(Tensor self, Scalar max) -> (Tensor)",
+                  [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+                    // Compute min(max(min_threshold, input), max_threshold)
+                    auto self = args[0].ITensorOrFreeze(ctx);
+                    auto clamp_layer_out = self;
+                    if (args[1].isIValue() && args[1].IValue()->isScalar()) {
+                      auto limit = args[1].unwrapToScalar().to<float>();
+                      clamp_layer_out = clamp_util(ctx, n, self, limit, nvinfer1::ElementWiseOperation::kMIN, "_min");
                     }
 
                     auto out = ctx->AssociateValueAndTensor(n->outputs()[0], clamp_layer_out);
-                    LOG_DEBUG("Clamp layer output tensor shape: " << out->getDimensions());
+                    LOG_DEBUG("clamp_max layer output tensor shape: " << out->getDimensions());
                     return true;
                   }})
         .pattern({"aten::sub.Tensor(Tensor self, Tensor other, Scalar alpha=1) -> "

tests/core/conversion/converters/test_element_wise.cpp

@@ -250,7 +250,7 @@ TEST(Converters, ATenRsubWithScalarConvertsCorrectly) {
 TEST(Converters, ATenClampMinConvertsCorrectly) {
   const auto graph = R"IR(
   graph(%x.1 : Tensor):
-          %2 : int = prim::Constant[value=-2]()
+          %2 : float = prim::Constant[value=1.5]()
           %3 : None = prim::Constant()
           %4 : Tensor = aten::clamp(%x.1, %2, %3)
           return (%4))IR";
@@ -260,7 +260,7 @@ TEST(Converters, ATenClampMinConvertsCorrectly) {
 TEST(Converters, ATenClampMaxConvertsCorrectly) {
   const auto graph = R"IR(
   graph(%x.1 : Tensor):
-          %2 : int = prim::Constant[value=3]()
+          %2 : float = prim::Constant[value=3.5]()
           %3 : None = prim::Constant()
           %4 : Tensor = aten::clamp(%x.1, %3, %2)
           return (%4))IR";
@@ -270,13 +270,31 @@ TEST(Converters, ATenClampMaxConvertsCorrectly) {
 TEST(Converters, ATenClampMinMaxConvertsCorrectly) {
   const auto graph = R"IR(
   graph(%x.1 : Tensor):
-          %2 : int = prim::Constant[value=3]()
-          %3 : int = prim::Constant[value=-2]()
+          %2 : float = prim::Constant[value=3.5]()
+          %3 : float = prim::Constant[value=1.5]()
           %4 : Tensor = aten::clamp(%x.1, %3, %2)
           return (%4))IR";
   pointwise_test_helper(graph, true);
 }
 
+TEST(Converters, ATenClampMinimumConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=2.5]()
+          %4 : Tensor = aten::clamp_min(%x.1, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
+TEST(Converters, ATenClampMaximumConvertsCorrectly) {
+  const auto graph = R"IR(
+  graph(%x.1 : Tensor):
+          %2 : float = prim::Constant[value=2.5]()
+          %4 : Tensor = aten::clamp_max(%x.1, %2)
+          return (%4))IR";
+  pointwise_test_helper(graph, true);
+}
+
 TEST(Converters, ATenGreaterThanConvertsCorrectly) {
   const auto graph = R"IR(
     graph(%0 : Tensor, %1 : Tensor):