add lookup_table op for Paddle-TRT (#54882)

Wangzheee · web-flow · commit 7c89b972f1a9 · 2023-06-29T11:24:22.000+08:00
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -2873,6 +2873,7 @@ USE_TRT_CONVERTER(fuse_eleadd_transpose)
 USE_TRT_CONVERTER(tanh_shrink)
 USE_TRT_CONVERTER(logsigmoid)
 USE_TRT_CONVERTER(lookup_table)
+USE_TRT_CONVERTER(lookup_table_v2)
 USE_TRT_CONVERTER(expand_v2)
 USE_TRT_CONVERTER(expand_as_v2)
 USE_TRT_CONVERTER(take_along_axis)
diff --git a/paddle/fluid/inference/tensorrt/convert/lookup_table_op.cc b/paddle/fluid/inference/tensorrt/convert/lookup_table_op.cc
@@ -24,8 +24,40 @@ class LookupTableOpConverter : public OpConverter {
                   const framework::Scope& scope,
                   bool test_mode) override {
     framework::OpDesc op_desc(op, nullptr);
-    VLOG(3)
-        << "convert lookup_table(lookup_table_v2) op to TensorRT IGatherLayer";
+    VLOG(3) << "convert lookup_table op to TensorRT IGatherLayer";
+
+    auto ids_name = op_desc.Input("Ids").front();
+    auto w_name = op_desc.Input("W").front();
+    auto out_name = op_desc.Output("Out").front();
+
+    auto* ids_tensor = engine_->GetITensor(ids_name);
+    auto* w_tensor = engine_->GetITensor(w_name);
+
+    std::vector<nvinfer1::ITensor*> after_shape_tensors;
+    // lookup_table'Ids has an additional one-dimensional dimension (*,1), need
+    // to reshape (*)
+    for (int i = 0; i < ids_tensor->getDimensions().nbDims - 1; ++i) {
+      after_shape_tensors.push_back(GetEleTensorOfShape(Shape(ids_tensor), i));
+    }
+
+    auto* reshape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *ids_tensor);
+    reshape_layer->setInput(1, *Concat(after_shape_tensors));
+    reshape_layer->setName(
+        ("reshape Ids for lookup_table(Output: " + out_name + ")").c_str());
+
+    auto* gather_layer = TRT_ENGINE_ADD_LAYER(
+        engine_, Gather, *w_tensor, *reshape_layer->getOutput(0), 0);
+    RreplenishLayerAndOutput(gather_layer, "gather", {out_name}, test_mode);
+  }
+};
+
+class LookupTableV2OpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    framework::OpDesc op_desc(op, nullptr);
+    VLOG(3) << "convert lookup_table_v2 op to TensorRT IGatherLayer";
 
     auto ids_name = op_desc.Input("Ids").front();
     auto w_name = op_desc.Input("W").front();
@@ -45,3 +77,4 @@ class LookupTableOpConverter : public OpConverter {
 }  // namespace paddle
 
 REGISTER_TRT_OP_CONVERTER(lookup_table, LookupTableOpConverter);
+REGISTER_TRT_OP_CONVERTER(lookup_table_v2, LookupTableV2OpConverter);
diff --git a/paddle/fluid/inference/tensorrt/convert/op_converter.h b/paddle/fluid/inference/tensorrt/convert/op_converter.h
@@ -140,14 +140,6 @@ class OpConverter {
               platform::errors::Unimplemented("no OpConverter for optype [%s]",
                                               op_desc.Type()));
         }
-        // lookup_table_v2 == lookup_table
-        if (op_desc.Type() == "lookup_table_v2") {
-          it = Registry<OpConverter>::Global().Lookup("lookup_table");
-          PADDLE_ENFORCE_NOT_NULL(
-              it,
-              platform::errors::Unimplemented("no OpConverter for optype [%s]",
-                                              op_desc.Type()));
-        }
         if (!it) {
           it = Registry<OpConverter>::Global().Lookup(op_desc.Type());
         }
diff --git a/test/ir/inference/test_trt_convert_lookup_table.py b/test/ir/inference/test_trt_convert_lookup_table.py
@@ -0,0 +1,154 @@
+# Copyright (c) 2023 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
+#
+#     https://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.
+
+import unittest
+from functools import partial
+from typing import Any, Dict, List
+
+import numpy as np
+from program_config import ProgramConfig, TensorConfig
+from trt_layer_auto_scan_test import TrtLayerAutoScanTest
+
+import paddle.inference as paddle_infer
+
+
+class TrtConvertLookupTableV2Test(TrtLayerAutoScanTest):
+    def sample_program_configs(self):
+        self.trt_param.workspace_size = 102400
+
+        def generate_input1(dims, attrs: List[Dict[str, Any]]):
+            if dims == 1:
+                return np.array([[32], [2], [19]]).astype(np.int64)
+            elif dims == 2:
+                return np.array([[[3], [16], [24]], [[6], [4], [47]]]).astype(
+                    np.int64
+                )
+            else:
+                return np.array(
+                    [
+                        [
+                            [[3], [16], [24]],
+                            [[30], [16], [14]],
+                            [[2], [6], [24]],
+                        ],
+                        [[[3], [26], [34]], [[3], [16], [24]], [[3], [6], [4]]],
+                        [
+                            [[3], [16], [24]],
+                            [[53], [16], [54]],
+                            [[30], [1], [24]],
+                        ],
+                    ]
+                ).astype(np.int64)
+
+        def generate_input2(dims, attrs: List[Dict[str, Any]]):
+            return np.random.uniform(-1, 1, [64, 4]).astype('float32')
+
+        for dims in [1, 2, 3]:
+            self.dims = dims
+
+            ops_config = [
+                {
+                    "op_type": "lookup_table",
+                    "op_inputs": {"Ids": ["indices"], "W": ["data"]},
+                    "op_outputs": {"Out": ["out_data"]},
+                    "op_attrs": {},
+                }
+            ]
+            ops = self.generate_op_config(ops_config)
+
+            program_config = ProgramConfig(
+                ops=ops,
+                weights={
+                    "data": TensorConfig(
+                        data_gen=partial(generate_input2, {}, {})
+                    )
+                },
+                inputs={
+                    "indices": TensorConfig(
+                        data_gen=partial(generate_input1, dims, {})
+                    )
+                },
+                outputs=["out_data"],
+            )
+
+            yield program_config
+
+    def sample_predictor_configs(
+        self, program_config
+    ) -> (paddle_infer.Config, List[int], float):
+        def generate_dynamic_shape(attrs):
+            if self.dims == 1:
+                self.dynamic_shape.min_input_shape = {
+                    "indices": [1, 1],
+                    "data": [64, 4],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "indices": [16, 1],
+                    "data": [64, 4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "indices": [8, 1],
+                    "data": [64, 4],
+                }
+            elif self.dims == 2:
+                self.dynamic_shape.min_input_shape = {
+                    "indices": [1, 1, 1],
+                    "data": [64, 4],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "indices": [16, 32, 1],
+                    "data": [64, 4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "indices": [2, 16, 1],
+                    "data": [64, 4],
+                }
+            else:
+                self.dynamic_shape.min_input_shape = {
+                    "indices": [1, 1, 1, 1],
+                    "data": [64, 4],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "indices": [16, 16, 16, 1],
+                    "data": [64, 4],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "indices": [2, 8, 8, 1],
+                    "data": [64, 4],
+                }
+
+        def generate_trt_nodes_num(attrs, dynamic_shape):
+            return 1, 2
+
+        attrs = [
+            program_config.ops[i].attrs for i in range(len(program_config.ops))
+        ]
+
+        # for dynamic_shape mode
+        generate_dynamic_shape(attrs)
+        self.trt_param.precision = paddle_infer.PrecisionType.Float32
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), 1e-5
+        self.trt_param.precision = paddle_infer.PrecisionType.Half
+        yield self.create_inference_config(), generate_trt_nodes_num(
+            attrs, True
+        ), (1e-3, 1e-3)
+
+    def test(self):
+        self.run_test()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/ir/inference/test_trt_convert_lookup_table_v2.py b/test/ir/inference/test_trt_convert_lookup_table_v2.py
@@ -25,9 +25,11 @@
 
 class TrtConvertLookupTableV2Test(TrtLayerAutoScanTest):
     def sample_program_configs(self):
+        self.trt_param.workspace_size = 102400
+
         def generate_input1(dims, attrs: List[Dict[str, Any]]):
             if dims == 1:
-                return np.array([32]).astype(np.int64)
+                return np.array([32, 2, 19]).astype(np.int64)
             elif dims == 2:
                 return np.array([[3, 16, 24], [6, 4, 47]]).astype(np.int64)
             else:
@@ -82,37 +84,37 @@ def generate_dynamic_shape(attrs):
                     "data": [64, 4],
                 }
                 self.dynamic_shape.max_input_shape = {
-                    "indices": [1],
+                    "indices": [16],
                     "data": [64, 4],
                 }
                 self.dynamic_shape.opt_input_shape = {
-                    "indices": [1],
+                    "indices": [8],
                     "data": [64, 4],
                 }
             elif self.dims == 2:
                 self.dynamic_shape.min_input_shape = {
-                    "indices": [2, 3],
+                    "indices": [1, 1],
                     "data": [64, 4],
                 }
                 self.dynamic_shape.max_input_shape = {
-                    "indices": [2, 3],
+                    "indices": [16, 32],
                     "data": [64, 4],
                 }
                 self.dynamic_shape.opt_input_shape = {
-                    "indices": [2, 3],
+                    "indices": [2, 16],
                     "data": [64, 4],
                 }
             else:
                 self.dynamic_shape.min_input_shape = {
-                    "indices": [3, 3, 3],
+                    "indices": [1, 1, 1],
                     "data": [64, 4],
                 }
                 self.dynamic_shape.max_input_shape = {
-                    "indices": [3, 3, 3],
+                    "indices": [16, 16, 16],
                     "data": [64, 4],
                 }
                 self.dynamic_shape.opt_input_shape = {
-                    "indices": [3, 3, 3],
+                    "indices": [2, 8, 8],
                     "data": [64, 4],
                 }
 
