Merge pull request #12598 from NHZlX/add_tensorrt_softmax

Add tensorrt softmax

paddle/fluid/inference/analysis/analyzer.cc

@@ -44,7 +44,7 @@ class DfgPassManagerImpl final : public DfgPassManager {
     if (FLAGS_inference_analysis_enable_tensorrt_subgraph_engine) {
       auto trt_teller = [&](const Node* node) {
         std::unordered_set<std::string> teller_set(
-            {"elementwise_add", "mul", "conv2d", "pool2d", "relu"});
+            {"elementwise_add", "mul", "conv2d", "pool2d", "relu", "softmax"});
         if (!node->IsFunction()) return false;
 
         const auto* func = static_cast<const Function*>(node);

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

@@ -1,7 +1,7 @@
 # Add TRT tests
 nv_library(tensorrt_converter
   SRCS mul_op.cc conv2d_op.cc fc_op.cc pool2d_op.cc elementwise_op.cc
-activation_op.cc
+activation_op.cc softmax_op.cc
   DEPS tensorrt_engine operator scope framework_proto op_registry)
 
 nv_test(test_op_converter SRCS test_op_converter.cc DEPS
@@ -21,3 +21,6 @@ nv_test(test_trt_pool2d_op SRCS test_pool2d_op.cc pool2d_op.cc
 
 nv_test(test_trt_elementwise_op SRCS test_elementwise_op.cc elementwise_op.cc
         DEPS ${FLUID_CORE_MODULES} tensorrt_engine elementwise_add_op SERIAL)
+
+nv_test(test_trt_softmax_op SRCS test_softmax_op.cc softmax_op.cc
+        DEPS ${FLUID_CORE_MODULES} tensorrt_engine softmax_op SERIAL)

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

@@ -0,0 +1,49 @@
+/* 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"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+/*
+ * SoftMaxOp, ISoftMaxLayer in TRT. This Layer doesn't has weights.
+ */
+class SoftMaxOpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope, bool test_mode) override {
+    VLOG(4)
+        << "convert a fluid softmax op to tensorrt softmax layer without bias";
+    framework::OpDesc op_desc(op, nullptr);
+    // Declare inputs
+    auto* input1 = engine_->GetITensor(op_desc.Input("X")[0]);
+    auto* layer = TRT_ENGINE_ADD_LAYER(engine_, SoftMax,
+                                       *const_cast<nvinfer1::ITensor*>(input1));
+
+    auto output_name = op_desc.Output("Out")[0];
+    engine_->SetITensor(output_name, layer->getOutput(0));
+    if (test_mode) {
+      engine_->DeclareOutput(output_name);
+    }
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+USE_OP(softmax);
+REGISTER_TRT_OP_CONVERTER(softmax, SoftMaxOpConverter);

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

@@ -0,0 +1,49 @@
+/* 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/framework/op_registry.h"
+#include "paddle/fluid/inference/tensorrt/convert/ut_helper.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+TEST(SoftMaxOpConverter, main) {
+  framework::Scope scope;
+  std::unordered_set<std::string> parameters;
+  TRTConvertValidation validator(8, parameters, scope, 1000);
+
+  std::vector<int> tensor_shape{8, 10};
+  validator.DeclInputVar("softmax-X", tensor_shape,
+                         nvinfer1::DimsCHW(10, 1, 1));
+  validator.DeclOutputVar("softmax-Out", nvinfer1::DimsCHW(10, 1, 1));
+
+  // Prepare Op description
+  framework::OpDesc desc;
+  desc.SetType("softmax");
+  desc.SetInput("X", {"softmax-X"});
+  desc.SetOutput("Out", {"softmax-Out"});
+
+  LOG(INFO) << "set OP";
+  validator.SetOp(*desc.Proto());
+  LOG(INFO) << "execute";
+
+  validator.Execute(3);
+}
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+USE_OP(softmax);

paddle/fluid/inference/tensorrt/convert/ut_helper.h

@@ -79,6 +79,12 @@ class TRTConvertValidation {
   }
 
   // Declare a Variable as input with random initialization.
+  void DeclInputVar(const std::string& name, const std::vector<int> tensor_dims,
+                    const nvinfer1::Dims& trt_dims) {
+    DeclVar(name, tensor_dims);
+    engine_->DeclareInput(name, nvinfer1::DataType::kFLOAT, trt_dims);
+  }
+
   void DeclInputVar(const std::string& name, const nvinfer1::Dims& dims) {
     DeclVar(name, dims);
     // Declare TRT inputs.
@@ -94,12 +100,18 @@ class TRTConvertValidation {
     DeclVar(name, dims);
   }
 
-  // Declare a variable in a fluid Scope.
-  void DeclVar(const std::string& name, const nvinfer1::Dims& dims,
-               bool is_param = false) {
+  void DeclVar(const std::string& name, const std::vector<int> dim_vec) {
     platform::CPUPlace place;
     platform::CPUDeviceContext ctx(place);
 
+    auto* x = scope_.Var(name);
+    auto* x_tensor = x->GetMutable<framework::LoDTensor>();
+    x_tensor->Resize(framework::make_ddim(dim_vec));
+    RandomizeTensor(x_tensor, place, ctx);
+  }
+  // Declare a variable in a fluid Scope.
+  void DeclVar(const std::string& name, const nvinfer1::Dims& dims,
+               bool is_param = false) {
     // Init Fluid tensor.
     std::vector<int> dim_vec(dims.d, dims.d + dims.nbDims);
     // There is no batchsize in ITensor's shape, but We should add it to
@@ -107,10 +119,8 @@ class TRTConvertValidation {
     // if_add_batch_ flag is true, add the max batchsize to dim_vec.
     if (is_param != true && if_add_batch_ == true)
       dim_vec.insert(dim_vec.begin(), max_batch_size_);
-    auto* x = scope_.Var(name);
-    auto* x_tensor = x->GetMutable<framework::LoDTensor>();
-    x_tensor->Resize(framework::make_ddim(dim_vec));
-    RandomizeTensor(x_tensor, place, ctx);
+
+    DeclVar(name, dim_vec);
   }
 
   void SetOp(const framework::proto::OpDesc& desc) {