Skip to content

[TF-TRT] Fix TopK Combined NMS #40062

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jan 13, 2021
Merged

[TF-TRT] Fix TopK Combined NMS #40062

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
79 changes: 38 additions & 41 deletions tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5771,7 +5771,7 @@ Status ConvertSquaredDifference(OpConverterParams* params) {
return Status::OK();
}

#if IS_TRT_VERSION_GE(5, 1, 0, 0)
#if IS_TRT_VERSION_GE(7, 1, 3, 0)
Status ConvertCombinedNMS(OpConverterParams* params) {
TF_RETURN_IF_ERROR(
CheckInputsWeights(*params, {{"boxes", false},
Expand All @@ -5793,13 +5793,22 @@ Status ConvertCombinedNMS(OpConverterParams* params) {
// Validate tensors and weights (also set some of the needed plugin fields)
const auto boxes_dims = boxes_tensor->getDimensions();
const auto scores_dims = scores_tensor->getDimensions();
if (boxes_dims.nbDims != 3) {
if (!params->use_implicit_batch &&
(!HasStaticShape(boxes_dims) || !HasStaticShape(scores_dims))) {
return errors::Unimplemented(
"TensorRT BatchedNMS Plugin requires input with static shape");
}
const int offset = params->use_implicit_batch ? 0 : 1;
if (boxes_dims.nbDims != 3 + offset) {
return errors::InvalidArgument(
"TensorRT BatchedNMS Plugin input boxes must be 3-D excluding batch ",
"TensorRT BatchedNMS Plugin input boxes must be 4-D including batch ",
node_def.name());
}
const int num_classes = scores_dims.d[1];
bool box_check = boxes_dims.d[1] == 1 || boxes_dims.d[1] == num_classes;
const int class_idx = 1 + offset;
const int num_classes = scores_dims.d[class_idx];
const int num_boxes = boxes_dims.d[0 + offset];
bool box_check =
boxes_dims.d[class_idx] == 1 || boxes_dims.d[class_idx] == num_classes;
if (!box_check) {
return errors::InvalidArgument(
"TensorRT BatchedNMS Plugin third dimension of boxes must be either 1 "
Expand Down Expand Up @@ -5848,23 +5857,32 @@ Status ConvertCombinedNMS(OpConverterParams* params) {

if (params->validation_only) return Status::OK();

// TF op CombinedNonMaxSuppression doesn't have the option of
// not normalizing coordinates.
// TRT op is_normalized=False treats input corrdinates as pixels and
// calculates width/height as (max - min + 1).
//
// TF op CombinedNonMaxSuppression doesn't care about the normalization and
// calculates width/height as (max-min).
//
// We set is_normalized = true to be consistent with TF IOU calculaton.
const bool is_normalized = true;
// Set plugin fields and the field collection

TFAttrs attrs(node_def);
bool share_location = (boxes_dims.d[1] == 1);
bool share_location = (boxes_dims.d[class_idx] == 1);
const bool pad_per_class = attrs.get<bool>("pad_per_class");
int top_k;
const bool clip_boxes = attrs.get<bool>("clip_boxes");
int keep_top_k = 0;
if (pad_per_class) {
top_k = std::min(max_size_per_class * num_classes, max_total_size);
keep_top_k = std::min(max_size_per_class * num_classes, max_total_size);
} else {
top_k = max_total_size;
keep_top_k = max_total_size;
}
const int keep_top_k = top_k;
// According to the batchedNMS plugin description we need to set top_k so that
// keep_top_k <= top_k
// https://github.com/NVIDIA/TensorRT/tree/master/plugin/batchedNMSPlugin
const int top_k = std::max(num_boxes, keep_top_k);
float score_thresh = *(static_cast<float*>(score_threshold.GetValues()));
const int background_id = -1;
nvinfer1::PluginField fields[8] = {
nvinfer1::PluginField fields[9] = {
nvinfer1::PluginField{"shareLocation", &share_location,
nvinfer1::PluginFieldType::kINT32, 1},
nvinfer1::PluginField{"backgroundLabelId", &background_id,
Expand All @@ -5881,8 +5899,9 @@ Status ConvertCombinedNMS(OpConverterParams* params) {
nvinfer1::PluginFieldType::kFLOAT32, 1},
nvinfer1::PluginField{"isNormalized", &is_normalized,
nvinfer1::PluginFieldType::kINT32, 1},
};
nvinfer1::PluginFieldCollection fc{8, fields};
nvinfer1::PluginField{"clipBoxes", &clip_boxes,
nvinfer1::PluginFieldType::kINT32, 1}};
nvinfer1::PluginFieldCollection fc{9, fields};

// Get plugin creator
auto creator =
Expand All @@ -5907,33 +5926,11 @@ Status ConvertCombinedNMS(OpConverterParams* params) {

// Set plugin outputs
nvinfer1::ITensor* output_nmsed_boxes = layer->getOutput(1);
#if IS_TRT_VERSION_GE(6, 0, 0, 0)

// TRT6 fixes (removes) the extra last dimension in CombinedNMS outputs
nvinfer1::ITensor* output_num_detections = layer->getOutput(0);
nvinfer1::ITensor* output_nmsed_scores = layer->getOutput(2);
nvinfer1::ITensor* output_nmsed_classes = layer->getOutput(3);
#else
nvinfer1::ITensor* output_num_detections = nullptr;
nvinfer1::ITensor* output_nmsed_scores = nullptr;
nvinfer1::ITensor* output_nmsed_classes = nullptr;

auto shrink_last_dim = [&](int output_index, nvinfer1::ITensor** out_tensor) {
nvinfer1::ITensor* in_tensor = layer->getOutput(output_index);
nvinfer1::Dims dims = in_tensor->getDimensions();
if (dims.d[dims.nbDims - 1] != 1) {
return errors::Internal("Expect last dims to be 1, for tensor ",
DebugString(*in_tensor));
}
--dims.nbDims;
TF_RETURN_IF_ERROR(PrepareTensorForShape(
params->converter, TRT_TensorOrWeights(in_tensor), dims,
/*validation_only=*/false, out_tensor, node_def, output_index));
return Status::OK();
};
TF_RETURN_IF_ERROR(shrink_last_dim(2, &output_nmsed_scores));
TF_RETURN_IF_ERROR(shrink_last_dim(3, &output_nmsed_classes));
TF_RETURN_IF_ERROR(shrink_last_dim(0, &output_num_detections));
#endif // IS_TRT_VERSION_GE(6, 0, 0, 0)

params->outputs->push_back(TRT_TensorOrWeights(output_nmsed_boxes));
params->outputs->push_back(TRT_TensorOrWeights(output_nmsed_scores));
Expand All @@ -5942,7 +5939,7 @@ Status ConvertCombinedNMS(OpConverterParams* params) {

return Status::OK();
}
#endif // IS_TRT_VERSION_GE(5, 1, 0, 0)
#endif // IS_TRT_VERSION_GE(7, 1, 3, 0)

#if IS_TRT_VERSION_GE(6, 0, 0, 0)
Status ConvertResize(OpConverterParams* params) {
Expand Down Expand Up @@ -6083,7 +6080,7 @@ static void RegisterValidatableOpConverters(
#if IS_TRT_VERSION_GE(5, 1, 2, 0)
(*registration)["ClipByValue"] = ConvertClipByValue;
#endif
#if IS_TRT_VERSION_GE(5, 1, 0, 0)
#if IS_TRT_VERSION_GE(7, 1, 3, 0)
(*registration)["CombinedNonMaxSuppression"] = ConvertCombinedNMS;
#endif
(*registration)["AddN"] = ConvertAddN;
Expand Down
Loading