Skip to content

Commit

Permalink
Add support for INT64 types in TensorRT constant layer calibration (#…
Browse files Browse the repository at this point in the history
…21101)

This PR is a duplicate of the
#21041
Create this PR in case the original one can't be updated for patch
release timeline.
  • Loading branch information
chilo-ms authored Jun 20, 2024
1 parent 6817b01 commit e737547
Showing 1 changed file with 10 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
Expand All @@ -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) {
Expand All @@ -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;
}
}
Expand Down

0 comments on commit e737547

Please sign in to comment.