Add LogSoftmax conversion

Fix Softmax comments


[TF-TRT] Move LogSoftmax to use OpConverterBase


Fix compiler errors


clang-format

Undo changes to convert_nodes.cc


Fix comments

Author: pavanimajety

tensorflow/compiler/tf2tensorrt/BUILD

@@ -677,6 +677,7 @@ tf_cuda_library(
         "convert/ops/einsum.cc",
         "convert/ops/fill_ops.cc",
         "convert/ops/like_ops.cc",
+        "convert/ops/log_softmax.cc",
         "convert/ops/quantization_ops.cc",
         "convert/ops/slice_ops.cc",
         "convert/ops/tile.cc",

tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc

@@ -4825,13 +4825,14 @@ TEST_P(OpConverter_FP32_FP16_Test, ConvertSoftmax) {
     std::vector<float> expected_values;
   };
   std::vector<TestParams> test_params = {
-      TestParams{{2, 3},
-                 {0.09003057, 0.24472848, 0.66524094, 0.09003057, 0.24472848,
-                  0.66524094}},
-      TestParams{{6, 1}, {1, 1, 1, 1, 1, 1}},  // works with std input
-      TestParams{{1, 6},  // this works with arange(1,7) input
-                 {0.00426978, 0.01160646, 0.03154963, 0.08576079, 0.23312202,
-                  0.6336913}},
+      TestParams{/*input_dims=*/{2, 3},
+                 /*expected_values=*/{0.09003057, 0.24472848, 0.66524094,
+                                      0.09003057, 0.24472848, 0.66524094}},
+      TestParams{/*input_dims=*/{6, 1},
+                 /*expected_values=*/{1, 1, 1, 1, 1, 1}}, // works w/ std input
+      TestParams{/*input_dims=*/{1, 6},  // this works w/ arange(1,7) input
+                /*expected_values=*/{0.00426978, 0.01160646, 0.03154963,
+                                     0.08576079, 0.23312202, 0.6336913}}
   };
   std::vector<float> input_values{1, 2, 3, 4, 5, 6};
   for (auto p : test_params) {
@@ -4842,6 +4843,37 @@ TEST_P(OpConverter_FP32_FP16_Test, ConvertSoftmax) {
   }
 }
 
+TEST_P(OpConverter_FP32_FP16_Test, ConvertLogSoftmax) {
+  // Get the NodeDef for LogSoftMax.
+  Scope s = Scope::NewRootScope();
+  auto input = ops::Placeholder(s.WithOpName("logits"), tf_type_);
+  auto logsoftmax = ops::LogSoftmax(s.WithOpName("my_logsoftmax"), input);
+  const NodeDef& node_def = logsoftmax.operation.node()->def();
+
+  struct TestParams {
+    std::vector<int> input_dims;
+    std::vector<float> expected_values;
+  };
+
+  std::vector<TestParams> test_params = {
+      TestParams{/*input_dims=*/{2, 3},
+                 /*expected_values=*/{-2.4076061, -1.407606, -0.40760604,
+                                      -2.4076061, -1.407606, -0.40760604}},
+      TestParams{/*input_dims=*/{1, 6},
+                 /*expected_values=*/{-5.4561934, -4.4561934, -3.4561934,
+                                      -2.4561934, -1.4561933, -0.45619333}},
+      TestParams{/*input_dims=*/{6, 1},
+                 /*expected_values=*/{0, 0, 0, 0, 0, 0}}
+  };
+  std::vector<float> input_values{1, 2, 3, 4, 5, 6};
+  for (auto p : test_params) {
+    Reset();
+    AddTestTensor("logits", p.input_dims, input_values);
+    TestOpConverter("my_logsoftmax", node_def, p.input_dims, Status::OK(),
+                    Status::OK(), ArrayFloatNear(p.expected_values, 1e-3));
+  }
+}
+
 TEST_P(OpConverter_FP32_Test, ConvertSqueeze) {
   const bool use_implicit_batch = (trt_mode_ == TrtTestMode::kImplicitBatch);
   // Get the NodeDef for Squeeze.

tensorflow/compiler/tf2tensorrt/convert/ops/log_softmax.cc

@@ -0,0 +1,94 @@
+#if GOOGLE_CUDA && GOOGLE_TENSORRT
+
+#include "tensorflow/compiler/tf2tensorrt/convert/convert_nodes.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/op_converter_registry.h"
+#include "tensorflow/compiler/tf2tensorrt/convert/ops/layer_utils.h"
+
+namespace tensorflow {
+namespace tensorrt {
+namespace convert {
+
+class ConvertLogSoftmax : public OpConverterBase<ConvertLogSoftmax> {
+ public:
+  explicit ConvertLogSoftmax(OpConverterParams *params)
+      : OpConverterBase<ConvertLogSoftmax>(params) {}
+
+  static constexpr std::array<DataType, 3> AllowedDataTypes() {
+    return {DataType::DT_FLOAT, DataType::DT_HALF};
+  }
+
+  static constexpr std::array<InputArgSpec, 1> InputSpec() {
+    return std::array<InputArgSpec, 1>{
+        InputArgSpec::Create("logits", TrtInputArg::kTensor)};
+  }
+
+  Status Validate() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    const auto &node_def = params.node_def;
+
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+    if (!num_trt_dims && params.use_implicit_batch) {
+      return errors::InvalidArgument(
+          "TensorRT LogSoftmax cannot apply on the batch dimension");
+    }
+
+    return Status::OK();
+  }
+
+  Status Convert() {
+    const auto &params = *this->params_;
+    const auto &inputs = params.inputs;
+    const auto &node_def = params.node_def;
+
+    // Perform LogSoftmax operation:
+    // `logsoftmax = logits - log(reduce_sum(exp(logits), axis))`
+
+    // Get the logits tensor.
+    ITensorProxyPtr logits_tensor = inputs.at(0).tensor();
+    const int num_trt_dims = logits_tensor->getDimensions().nbDims;
+
+    // Exponent of logits.
+    nvinfer1::IUnaryLayer *exp = params.converter->network()->addUnary(
+        *logits_tensor->trt_tensor(), nvinfer1::UnaryOperation::kEXP);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(exp, node_def.name());
+    params.converter->SetLayerName(exp, node_def, "exp");
+
+    // Reduce-sum operation across the final dimension.
+    nvinfer1::IReduceLayer *reduced_sum =
+        params.converter->network()->addReduce(
+            *exp->getOutput(0), nvinfer1::ReduceOperation::kSUM,
+            (1 << (num_trt_dims - 1)), /*Reduce across final dimension*/
+            true /*Keep reduced dims*/);
+    params.converter->SetLayerName(reduced_sum, node_def, "reduced_sum");
+
+    // Logarithm of reduced_sum.
+    nvinfer1::IUnaryLayer *log_reduced_sum =
+        params.converter->network()->addUnary(*reduced_sum->getOutput(0),
+                                              nvinfer1::UnaryOperation::kLOG);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(log_reduced_sum, node_def.name());
+    params.converter->SetLayerName(log_reduced_sum, node_def,
+                                   "log_reduced_sum");
+
+    // Finally, get the output by subtracting log_reduced_sum from logits.
+    nvinfer1::IElementWiseLayer *sub =
+        params.converter->network()->addElementWise(
+            *logits_tensor->trt_tensor(), *log_reduced_sum->getOutput(0),
+            nvinfer1::ElementWiseOperation::kSUB);
+    TFTRT_RETURN_ERROR_IF_NULLPTR(sub, node_def.name());
+    params.converter->SetLayerName(sub, node_def, "sub");
+
+    params.outputs->push_back(TRT_TensorOrWeights(sub->getOutput(0)));
+    return Status::OK();
+  }
+};
+
+REGISTER_DEFAULT_TRT_OP_CONVERTER(MakeConverterFunction<ConvertLogSoftmax>(),
+                                  "LogSoftmax");
+
+}  // namespace convert
+}  // namespace tensorrt
+}  // namespace tensorflow
+#endif  // GOOGLE_CUDA && GOOGLE_TENSORRT