未验证 提交 5e880840 编写于 作者: W wenbin 提交者: GitHub

correct slice serialize data (#36588)

* slice

* add UT
上级 1962d3af
......@@ -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.
......
......@@ -86,5 +86,19 @@ class SlicePluginTRTTestFp16(SlicePluginTRTTest):
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()
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册