Implement leaky relu tensorRT converter

paddle/fluid/inference/analysis/passes/ir_analysis_compose_pass.cc

@@ -45,7 +45,8 @@ void IrAnalysisComposePass::InitTensorRTAttrs(Argument *argument) {
       std::unordered_set<std::string> teller_set(
           {"mul", "conv2d", "pool2d", "relu", "softmax", "sigmoid",
            "depthwise_conv2d", "batch_norm", "concat", "tanh", "pad",
-           "elementwise_add", "dropout", "split", "prelu", "conv2d_transpose"});
+           "elementwise_add", "dropout", "split", "prelu", "conv2d_transpose",
+           "leaky_relu"});
       if (!node->IsOp()) return false;
 
       if (teller_set.count(node->Op()->Type())) {

paddle/fluid/inference/api/analysis_predictor.cc

@@ -551,4 +551,5 @@ USE_TRT_CONVERTER(pad);
 USE_TRT_CONVERTER(split);
 USE_TRT_CONVERTER(prelu);
 USE_TRT_CONVERTER(conv2d_transpose);
+USE_TRT_CONVERTER(leaky_relu);
 #endif

paddle/fluid/inference/tensorrt/convert/CMakeLists.txt

@@ -2,7 +2,7 @@
 nv_library(tensorrt_converter
   SRCS mul_op.cc conv2d_op.cc fc_op.cc pool2d_op.cc elementwise_op.cc
 batch_norm_op.cc activation_op.cc softmax_op.cc concat_op.cc dropout_op.cc
-pad_op.cc split_op.cc prelu_op.cc
+pad_op.cc split_op.cc prelu_op.cc leaky_relu_op.cc
   DEPS tensorrt_engine tensorrt_plugin operator scope framework_proto op_registry)
 
 nv_test(test_op_converter SRCS test_op_converter.cc DEPS
@@ -37,3 +37,5 @@ nv_test(test_trt_split_op SRCS test_split_op.cc split_op.cc
 nv_test(test_trt_prelu_op SRCS test_prelu_op.cc prelu_op.cc
         DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine tensorrt_plugin
         prelu_op SERIAL)
+nv_test(test_trt_leaky_relu_op SRCS test_leaky_relu_op.cc leaky_relu_op.cc
+        DEPS ${FLUID_CORE_MODULES} ${GLOB_OPERATOR_DEPS} tensorrt_engine activation_op SERIAL)

paddle/fluid/inference/tensorrt/convert/leaky_relu_op.cc

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+#include "paddle/fluid/inference/tensorrt/plugin/prelu_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+// LeakyRelu converter from fluid to tensorRT
+class LeakyReluOpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope, bool test_mode) override {
+    VLOG(4) << "convert fluid leaky_relu op to tensorrt layer";
+
+    framework::OpDesc op_desc(op, nullptr);
+    // Declare inputs
+    int input_num = op_desc.Input("X").size();
+    PADDLE_ENFORCE(input_num == 1);
+    auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
+    // Get output
+    size_t output_num = op_desc.Output("Out").size();
+    PADDLE_ENFORCE(output_num == 1);
+    // Get attrs
+    float alpha = boost::get<float>(op_desc.GetAttr("alpha"));
+
+    platform::CPUPlace place;
+    std::unique_ptr<framework::LoDTensor> alpha_tensor(
+        new framework::LoDTensor());
+    alpha_tensor->Resize(framework::make_ddim({2}));
+    float* alpha_data = alpha_tensor->mutable_data<float>(place);
+    alpha_data[0] = alpha;
+    alpha_data[1] = 1.f - alpha;
+    // the leaky relu formula y = (x > 0) ? x : alpha * x is equal to
+    // y = alpha * x + (x > 0) ? (1 - alpha) * x : 0
+    TensorRTEngine::Weight scale{nvinfer1::DataType::kFLOAT, &alpha_data[0], 1};
+    TensorRTEngine::Weight shift{nvinfer1::DataType::kFLOAT, nullptr, 0};
+    TensorRTEngine::Weight power{nvinfer1::DataType::kFLOAT, nullptr, 0};
+    // y_scale = alpha * x
+    auto* scale_layer = TRT_ENGINE_ADD_LAYER(
+        engine_, Scale, *input, nvinfer1::ScaleMode::kUNIFORM, shift.get(),
+        scale.get(), power.get());
+    PADDLE_ENFORCE(nullptr != scale_layer);
+    // y_relu = (x > 0) : x : 0
+    auto* relu_layer = TRT_ENGINE_ADD_LAYER(engine_, Activation, *input,
+                                            nvinfer1::ActivationType::kRELU);
+    PADDLE_ENFORCE(nullptr != relu_layer);
+    //
+    TensorRTEngine::Weight sub_scale{nvinfer1::DataType::kFLOAT, &alpha_data[1],
+                                     1};
+    auto* scale_relu_layer =
+        TRT_ENGINE_ADD_LAYER(engine_, Scale, *(relu_layer->getOutput(0)),
+                             nvinfer1::ScaleMode::kUNIFORM, shift.get(),
+                             sub_scale.get(), power.get());
+    PADDLE_ENFORCE(nullptr != scale_relu_layer);
+    auto* output_layer =
+        TRT_ENGINE_ADD_LAYER(engine_, ElementWise, *(scale_layer->getOutput(0)),
+                             *(scale_relu_layer->getOutput(0)),
+                             nvinfer1::ElementWiseOperation::kSUM);
+    PADDLE_ENFORCE(nullptr != output_layer);
+    // keep alpha tensor to avoid release it's memory
+    engine_->weight_map[op_desc.Input("alpha")[0]] = std::move(alpha_tensor);
+
+    std::string layer_name = "leaky_relu (Output: ";
+    auto output_name = op_desc.Output("Out")[0];
+    output_layer->getOutput(0)->setName(output_name.c_str());
+    engine_->SetITensor(output_name, output_layer->getOutput(0));
+    layer_name += output_name;
+    if (test_mode) {
+      engine_->DeclareOutput(output_name);
+    }
+    output_layer->setName((layer_name + ")").c_str());
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(leaky_relu, LeakyReluOpConverter);

paddle/fluid/inference/tensorrt/convert/test_leaky_relu_op.cc

+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <gtest/gtest.h>
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+#include "paddle/fluid/inference/tensorrt/convert/ut_helper.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+TEST(leaky_relu_op, test_channel_wise) {
+  std::unordered_set<std::string> parameters({"leaky_relu_alpha"});
+  framework::Scope scope;
+  TRTConvertValidation validator(10, parameters, scope, 1000);
+  validator.DeclInputVar("leaky_relu_input", nvinfer1::DimsCHW(3, 2, 2));
+  validator.DeclOutputVar("leaky_relu_out", nvinfer1::DimsCHW(3, 2, 2));
+
+  // Prepare Op description
+  framework::OpDesc desc;
+  desc.SetType("leaky_relu");
+  desc.SetInput("X", {"leaky_relu_input"});
+  desc.SetOutput("Out", {"leaky_relu_out"});
+
+  desc.SetAttr("alpha", 0.1f);
+
+  validator.SetOp(*desc.Proto());
+
+  validator.Execute(1);
+}
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+// USE_OP(leaky_relu);
+USE_OP(leaky_relu);

paddle/fluid/inference/tensorrt/plugin/CMakeLists.txt

-nv_library(tensorrt_plugin SRCS trt_plugin.cc split_op_plugin.cu prelu_op_plugin.cu DEPS enforce device_context)
+nv_library(tensorrt_plugin SRCS trt_plugin.cc split_op_plugin.cu prelu_op_plugin.cu DEPS enforce tensorrt_engine)