correct slice serialize data (#36588)

* slice * add UT
PaddlePaddle · Oct 22, 2021 · 5e88084 · 5e88084
1 parent 1962d3a
commit 5e88084
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/paddle/fluid/inference/tensorrt/plugin/slice_op_plugin.cu b/paddle/fluid/inference/tensorrt/plugin/slice_op_plugin.cu
@@ -65,6 +65,7 @@ SlicePlugin::SlicePlugin(void const *serial_data, size_t serial_length) {
   DeserializeValue(&serial_data, &serial_length, &starts_);
   DeserializeValue(&serial_data, &serial_length, &ends_);
   DeserializeValue(&serial_data, &serial_length, &axes_);
+  DeserializeValue(&serial_data, &serial_length, &with_fp16_);
   cudaEventCreate(&copy_event_);
   cudaStreamCreate(&copy_stream_);
 }
@@ -187,17 +188,17 @@ int SlicePlugin::enqueue(int batch_size, const void *const *inputs,
 }
 
 size_t SlicePlugin::getSerializationSize() const TRT_NOEXCEPT {
-  return getBaseSerializationSize() + SerializedSize(getPluginType()) +
-         SerializedSize(starts_) + SerializedSize(ends_) +
-         SerializedSize(axes_);
+  return getBaseSerializationSize() + SerializedSize(starts_) +
+         SerializedSize(ends_) + SerializedSize(axes_) +
+         SerializedSize(with_fp16_);
 }
 
 void SlicePlugin::serialize(void *buffer) const TRT_NOEXCEPT {
-  SerializeValue(&buffer, getPluginType());
   serializeBase(buffer);
   SerializeValue(&buffer, starts_);
   SerializeValue(&buffer, ends_);
   SerializeValue(&buffer, axes_);
+  SerializeValue(&buffer, with_fp16_);
 }
 
 // Dynamic Plugin below.

diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_slice_plugin.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_slice_plugin.py
@@ -86,5 +86,19 @@ def setUpTensorRTParams(self):
         self.enable_trt = True
 
 
+class StaticSlicePluginTRTTestFp16(SlicePluginTRTTest):
+    def setUpTensorRTParams(self):
+        self.trt_parameters = SlicePluginTRTTest.TensorRTParam(
+            1 << 30, 32, 1, AnalysisConfig.Precision.Half, True, False)
+        self.enable_trt = True
+
+
+class StaticSlicePluginTRTTestFp32(SlicePluginTRTTest):
+    def setUpTensorRTParams(self):
+        self.trt_parameters = SlicePluginTRTTest.TensorRTParam(
+            1 << 30, 32, 1, AnalysisConfig.Precision.Float32, True, False)
+        self.enable_trt = True
+
+
 if __name__ == "__main__":
     unittest.main()