[TF FE] Support NonMaxSuppression with named outputs (#16835)

* [TF FE] Support NonMaxSuppression with named outputs

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>

* Simplify the test for NMS named outputs

* Share a script for test model generation

---------

Signed-off-by: Kazantsev, Roman <roman.kazantsev@intel.com>
This commit is contained in:
Roman Kazantsev 2023-04-11 21:14:59 +04:00 committed by GitHub
parent 7513e9dee1
commit 9e89b6c5f6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 121 additions and 7 deletions

View File

@ -563,3 +563,71 @@ TEST_F(TransformationTestsF, ResourceGatherModel) {
model_ref = make_shared<Model>(OutputVector{mul}, ParameterVector{ind1, ind2});
}
}
TEST_F(TransformationTestsF, NonMaxSuppressionWithNamedOutputs) {
// The purpose of this test is to check that named output ports of TensorFlow NMS operation are connected correctly
// to its consumers
{ model = convert_model("nms_named_outputs/nms_named_outputs.pb"); }
{
// prepare the first input for NMS
auto boxes = make_shared<Parameter>(f32, PartialShape{2, 4});
auto const_zero = make_shared<Constant>(i32, Shape{1}, 0);
auto unsqueeze = make_shared<Unsqueeze>(boxes, const_zero);
// prepare the second input for NMS
auto scores = make_shared<Parameter>(f32, PartialShape{2});
auto const_one_zero = make_shared<Constant>(i32, Shape{2}, vector<int32_t>{0, 1});
auto unsqueeze_2 = make_shared<Unsqueeze>(scores, const_one_zero);
// create NMS node
auto max_output_size = make_shared<Constant>(i32, Shape{}, 50);
auto iou_threshold = make_shared<Constant>(f32, Shape{}, 0.4f);
auto score_threshold = make_shared<Constant>(f32, Shape{}, 0.3f);
auto soft_nms_sigma = make_shared<Constant>(f32, Shape{}, 0.1f);
auto nms = make_shared<NonMaxSuppression>(unsqueeze,
unsqueeze_2,
max_output_size,
iou_threshold,
score_threshold,
soft_nms_sigma,
NonMaxSuppression::BoxEncodingType::CORNER,
false,
i32);
// compute the first output - selected_indices
auto slice_const_one = make_shared<Constant>(i32, Shape{1}, 1);
auto slice_const_one_2 = make_shared<Constant>(i32, Shape{1}, 1);
auto slice_const_two = make_shared<Constant>(i32, Shape{1}, 2);
auto slice_const_three = make_shared<Constant>(i32, Shape{1}, 3);
auto slice =
make_shared<Slice>(nms->output(0), slice_const_two, slice_const_three, slice_const_one, slice_const_one_2);
Output<Node> selected_indices = make_shared<Squeeze>(slice, slice_const_one_2);
// compute the second output - selected_scores
auto slice2_const_one = make_shared<Constant>(i32, Shape{1}, 1);
auto slice2_const_one_2 = make_shared<Constant>(i32, Shape{1}, 1);
auto slice2_const_two = make_shared<Constant>(i32, Shape{1}, 2);
auto slice2_const_three = make_shared<Constant>(i32, Shape{1}, 3);
auto slice2 =
make_shared<Slice>(nms->output(1), slice_const_two, slice_const_three, slice_const_one, slice_const_one_2);
Output<Node> selected_scores = make_shared<Squeeze>(slice2, slice_const_one_2);
selected_scores = make_shared<ConvertLike>(selected_scores, boxes);
selected_scores = make_shared<Convert>(selected_scores, i32);
// compute the third output - valid_outputs
Output<Node> valid_outputs = make_shared<Squeeze>(nms->output(2));
// make post-processing before the concatenation
auto const_minus_one = make_shared<Constant>(i32, Shape{1}, -1);
selected_indices = make_shared<Reshape>(selected_indices, const_minus_one, false);
auto const_minus_one_2 = make_shared<Constant>(i32, Shape{1}, -1);
selected_scores = make_shared<Reshape>(selected_scores, const_minus_one_2, false);
auto const_minus_one_3 = make_shared<Constant>(i32, Shape{1}, -1);
valid_outputs = make_shared<Reshape>(valid_outputs, const_minus_one_3, false);
// concatenate all outputs in order to have the single output
auto concat = make_shared<Concat>(OutputVector{selected_indices, selected_scores, valid_outputs}, 0);
model_ref = make_shared<Model>(OutputVector{concat}, ParameterVector{boxes, scores});
}
}

View File

@ -0,0 +1,45 @@
# Copyright (C) 2018-2023 Intel Corporation
# SPDX-License-Identifier: Apache-2.0
import os
import sys
import tensorflow.compat.v1 as tf
def main():
tf.compat.v1.reset_default_graph()
@tf.function
def second_func(boxes, scores):
# the second function is used to obtain the body graph with NonMaxSuppressionV5
# only body graphs use named output ports
selected_indices, selected_scores, valid_outputs = tf.raw_ops.NonMaxSuppressionV5(boxes=boxes, scores=scores,
max_output_size=50,
iou_threshold=0.4,
score_threshold=0.3,
soft_nms_sigma=0.1)
return selected_indices, selected_scores, valid_outputs
@tf.function
def first_func(boxes, scores):
selected_indices, selected_scores, valid_outputs = second_func(boxes, scores)
# the post-processing part of the test model to get the single output
selected_indices = tf.raw_ops.Reshape(tensor=selected_indices, shape=[-1])
selected_scores = tf.raw_ops.Cast(x=selected_scores, DstT=tf.int32)
selected_scores = tf.raw_ops.Reshape(tensor=selected_scores, shape=[-1])
valid_outputs = tf.raw_ops.Reshape(tensor=valid_outputs, shape=[-1])
concat = tf.raw_ops.Concat(concat_dim=0, values=[selected_indices, selected_scores, valid_outputs])
return concat
tf_net = first_func.get_concrete_function(
tf.constant([[0.1, 0.2, 0.3, 0.5], [0.4, 0.2, 0.1, 0.8]], dtype=tf.float32),
tf.constant([0, 2], dtype=tf.float32)).graph.as_graph_def()
tf.io.write_graph(tf_net, os.path.join(sys.argv[1], "nms_named_outputs"),
"nms_named_outputs.pb", False)
if __name__ == "__main__":
main()

View File

@ -86,7 +86,7 @@ OP_CONVERTER(translate_mat_mul_op);
OP_CONVERTER(translate_matrix_diag_op);
OP_CONVERTER(translate_max_pool_op);
OP_CONVERTER(translate_mirror_pad_op);
OP_CONVERTER(translate_non_max_suppression_op);
OP_CONVERTER_NAMED(translate_non_max_suppression_op);
OP_CONVERTER(translate_parallel_dynamic_stitch_op);
OP_CONVERTER(translate_placeholder_op);
OP_CONVERTER(translate_placeholder_with_default_op);

View File

@ -40,7 +40,7 @@ Output<Node> normalize_selected_indices(const Output<Node>& ov_selected_indices,
return selected_indices;
}
OutputVector translate_non_max_suppression_op(const NodeContext& node) {
NamedOutputVector translate_non_max_suppression_op(const NodeContext& node) {
default_op_checks(node,
3,
{"NonMaxSuppression",
@ -65,7 +65,6 @@ OutputVector translate_non_max_suppression_op(const NodeContext& node) {
auto ov_scores = make_shared<Unsqueeze>(scores, scores_axes);
const auto& op_type = node.get_op_type();
OutputVector results;
// set all thresholds to zero, default values
Output<Node> iou_threshold = make_shared<Constant>(element::f32, Shape{}, 0.0);
@ -124,21 +123,23 @@ OutputVector translate_non_max_suppression_op(const NodeContext& node) {
element::i32);
auto tf_selected_indices =
normalize_selected_indices(non_max_suppression->output(0), max_output_size, pad_to_max_output_size);
results.push_back(tf_selected_indices);
NamedOutputVector named_results;
named_results.push_back({"selected_indices", tf_selected_indices});
Output<Node> tf_selected_scores;
if (selected_scores) {
tf_selected_scores =
normalize_selected_indices(non_max_suppression->output(1), max_output_size, pad_to_max_output_size);
tf_selected_scores = make_shared<ConvertLike>(tf_selected_scores, boxes)->output(0);
results.push_back(tf_selected_scores);
named_results.push_back({"selected_scores", tf_selected_scores});
}
Output<Node> tf_valid_outputs;
if (valid_outputs) {
// In TensorFlow valid_outputs output is a scalar
tf_valid_outputs = make_shared<Squeeze>(non_max_suppression->output(2))->output(0);
results.push_back(tf_valid_outputs);
named_results.push_back({"valid_outputs", tf_valid_outputs});
}
// set output tensor names based on a number of outputs
@ -152,7 +153,7 @@ OutputVector translate_non_max_suppression_op(const NodeContext& node) {
set_node_name(node.get_name() + ":2", tf_valid_outputs.get_node_shared_ptr());
}
return results;
return named_results;
}
} // namespace op
} // namespace tensorflow