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

Add support for INT64 types in TensorRT constant layer calibration (#21101)

Browse source 
This PR is a duplicate of the
https://github.com/microsoft/onnxruntime/pull/21041
Create this PR in case the original one can't be updated for patch
release timeline.
This commit is contained in:
Chi Lo 2024-06-19 18:36:26 -07:00 committed by GitHub
parent 6817b013b9
commit e737547862
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 10 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -71,6 +71,7 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
auto dynamic_range_iter = dynamic_range_map.find(tensor_name);
if (dynamic_range_iter != dynamic_range_map.end()) {
if (!network.getInput(i)->setDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) {
LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for network input " << tensor_name;
return false;
}
}
@ -84,10 +85,12 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
auto dynamic_range_iter = dynamic_range_map.find(tensor_name);
if (dynamic_range_iter != dynamic_range_map.end()) {
if (!trt_layer->getOutput(j)->setDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) {
LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for tensor " << tensor_name;
return false;
}
} else if (trt_layer->getType() == nvinfer1::LayerType::kCONSTANT) {
nvinfer1::IConstantLayer* const_layer = static_cast<nvinfer1::IConstantLayer*>(trt_layer);
const std::string const_layer_name = const_layer->getName();
auto trt_weights = const_layer->getWeights();
double max_weight = std::numeric_limits<double>::min();
for (int64_t k = 0, end = trt_weights.count; k < end; ++k) {
@ -108,13 +111,19 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
case nvinfer1::DataType::kINT32:
weight = static_cast<const int32_t*>(trt_weights.values)[k];
break;
#if NV_TENSORRT_MAJOR >= 10
case nvinfer1::DataType::kINT64:
weight = static_cast<double>(static_cast<const int64_t*>(trt_weights.values)[k]);
break;
#endif // NV_TENSORRT_MAJOR >= 10
default:
LOGS_DEFAULT(ERROR) << "Found unsupported datatype!";
LOGS_DEFAULT(ERROR) << "Found unsupported datatype for layer " << const_layer_name;
return false;
}
max_weight = std::max(max_weight, std::abs(weight));
}
if (!trt_layer->getOutput(j)->setDynamicRange(static_cast<float>(-max_weight), static_cast<float>(max_weight))) {
LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for layer " << const_layer_name;
return false;
}
}