saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-11 17:48:34 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add TRT uint8 support (#12570)

Browse source 
* uint8 support

Signed-off-by: Kevin Chen <kevinch@nvidia.com>

* Handle outputs as well

Signed-off-by: Kevin Chen <kevinch@nvidia.com>

Signed-off-by: Kevin Chen <kevinch@nvidia.com>
This commit is contained in:
Kevin Chen 2022-08-15 08:22:50 -07:00 committed by GitHub
parent 95f2a3e7e0
commit 73da3f3705
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 20 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file

@ -1737,6 +1737,16 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
}
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: {
auto input_tensor_ptr = ort.GetTensorData<uint8_t>(input_tensor);
if (input_tensor_ptr == nullptr) {
scratch_buffers.push_back(IAllocator::MakeUniquePtr<void>(alloc, sizeof(uint8_t)));
buffers[binding_index] = scratch_buffers.back().get();
} else {
buffers[binding_index] = const_cast<uint8_t*>(input_tensor_ptr);
}
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: {
auto input_tensor_ptr = ort.GetTensorData<int32_t>(input_tensor);
if (input_tensor_ptr == nullptr) {
@ -1870,6 +1880,16 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
}
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: {
auto output_tensor_ptr = ort.GetTensorMutableData<uint8_t>(output_tensor[i]);
if (output_tensor_ptr == nullptr) {
scratch_buffers.push_back(IAllocator::MakeUniquePtr<void>(alloc, sizeof(uint8_t)));
buffers[binding_index] = scratch_buffers.back().get();
} else {
buffers[binding_index] = output_tensor_ptr;
}
break;
}
case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: {
auto output_tensor_ptr = ort.GetTensorMutableData<int32_t>(output_tensor[i]);
if (output_tensor_ptr == nullptr) {