Auto batch lost label fix (#19535)

* Restored opset1::Reshape label peropagation for -1 special value

* Lets opset1::Reshape keep same shape infer. Makes FindBatch transformation keep labels in output shapes of Result node

* uses Parameter from correct namespace
This commit is contained in:
Evgenya Stepyreva 2023-09-04 19:11:09 +04:00 committed by GitHub
parent 15685e0141
commit fef4d4d641
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 102 additions and 35 deletions

View File

@ -45,8 +45,10 @@ void mark_layout_independent_batch(const std::shared_ptr<ov::opset1::Parameter>&
P2Btype& map);
void mark_with_unique_dimension_labels(const std::shared_ptr<Model>& m, const ov::DimensionTracker& dt);
void restore_original_dimensions(
const std::shared_ptr<ov::Model>& model,
const std::map<std::shared_ptr<ov::opset1::Parameter>, ov::PartialShape>& parameter_to_shape,
bool leave_batch_dynamic = true);
bool leave_batch_dynamic = true,
bool clear_labels = false);
bool check_batch_tracks_through_all_the_nodes(const std::shared_ptr<ov::Model>& m);
P2Btype find_batch(const std::shared_ptr<ov::Model>& m);
bool detach_detection_output(const std::shared_ptr<ov::Model>& f);

View File

@ -21,15 +21,16 @@
#include "openvino/op/result.hpp"
#include "openvino/op/shape_of.hpp"
void ov::batch_util::mark_with_unique_dimension_labels(const std::shared_ptr<ov::Model>& f,
void ov::batch_util::mark_with_unique_dimension_labels(const std::shared_ptr<ov::Model>& m,
const ov::DimensionTracker& dt) {
ov::label_t i = 1;
for (auto& parameter : f->get_parameters()) {
for (auto& parameter : m->get_parameters()) {
ov::PartialShape new_shape = ov::PartialShape::dynamic(parameter->get_partial_shape().rank());
for (auto& dim : new_shape)
dt.set_up_for_tracking(dim, i++);
parameter->set_partial_shape(new_shape);
}
m->validate_nodes_and_infer_types();
}
void ov::batch_util::mark_batch(const std::shared_ptr<ov::op::v0::Parameter>& parameter,
@ -161,15 +162,17 @@ P2Btype ov::batch_util::find_batch(const std::shared_ptr<ov::Model>& f) {
for (auto& result : layout_independent_results)
// there are no layout obvious operations on the Parameter-Result path
// considering the outer-most matching dimension is batch
// considering the outermost matching dimension is batch
mark_layout_independent_batch(parameter, result->shared_from_this(), parameter_to_batch_labels);
}
return parameter_to_batch_labels;
}
void ov::batch_util::restore_original_dimensions(
const std::shared_ptr<ov::Model>& model,
const std::map<std::shared_ptr<ov::op::v0::Parameter>, ov::PartialShape>& parameter_to_shape,
bool leave_batch_dynamic) {
bool leave_batch_dynamic,
bool clear_labels) {
for (const auto& item : parameter_to_shape) {
const auto& batch_marked_shape = item.first->get_partial_shape();
auto original_shape = item.second;
@ -180,11 +183,34 @@ void ov::batch_util::restore_original_dimensions(
if (const auto& label = ov::DimensionTracker::get_label(batch_marked_shape[n])) {
if (leave_batch_dynamic)
original_shape[n] = Dimension::dynamic();
ov::DimensionTracker::set_label(original_shape[n], label);
if (!clear_labels)
ov::DimensionTracker::set_label(original_shape[n], label);
}
}
item.first->set_partial_shape(original_shape);
}
std::unordered_map<std::shared_ptr<ov::op::v0::Result>, ov::PartialShape> output_to_shape;
if (!clear_labels) {
for (const auto& result : model->get_results())
output_to_shape[result] = result->get_output_partial_shape(0);
}
model->validate_nodes_and_infer_types();
if (!clear_labels) {
for (const auto& item : output_to_shape) {
auto labeled_shape = item.second, current_shape = item.first->get_output_partial_shape(0);
auto labeled_rank = labeled_shape.rank(), current_rank = current_shape.rank();
if (labeled_rank.is_static() && current_rank.is_static() && labeled_rank == current_rank) {
for (size_t i = 0; i < labeled_shape.size(); ++i) {
auto label = ov::DimensionTracker::get_label(labeled_shape[i]);
if (label != ov::no_label)
ov::DimensionTracker::set_label(current_shape[i], label);
}
item.first->set_output_type(0, item.first->get_element_type(), current_shape);
}
}
}
}
bool ov::batch_util::check_batch_tracks_through_all_the_nodes(const std::shared_ptr<ov::Model>& f) {
@ -252,6 +278,19 @@ bool ov::batch_util::detach_detection_output(const std::shared_ptr<ov::Model>& f
return !new_outputs.empty() || !outputs_to_delete.empty();
}
std::map<std::shared_ptr<ov::op::v0::Parameter>, ov::PartialShape> collect_original_input_shapes(
const std::shared_ptr<ov::Model>& m) {
const auto& parameters = m->get_parameters();
std::map<std::shared_ptr<ov::op::v0::Parameter>, ov::PartialShape> parameter_to_shape;
for (const auto& parameter : parameters) {
auto shape = parameter->get_partial_shape();
if (shape.rank().is_dynamic())
return {};
parameter_to_shape[parameter] = shape;
}
return parameter_to_shape;
}
bool ov::pass::FindBatch::run_on_model(const std::shared_ptr<ov::Model>& m) {
RUN_ON_MODEL_SCOPE(FindBatch);
auto te = std::make_shared<ov::TableOfEquivalence>();
@ -261,38 +300,20 @@ bool ov::pass::FindBatch::run_on_model(const std::shared_ptr<ov::Model>& m) {
if (detach_do)
model_has_changed |= batch_util::detach_detection_output(m);
const auto& parameters = m->get_parameters();
std::map<std::shared_ptr<ov::op::v0::Parameter>, PartialShape> parameter_to_shape;
for (const auto& parameter : parameters) {
auto shape = parameter->get_partial_shape();
if (shape.rank().is_dynamic())
return model_has_changed;
parameter_to_shape[parameter] = shape;
}
auto parameter_to_shape = collect_original_input_shapes(m);
if (parameter_to_shape.empty())
return model_has_changed;
ov::batch_util::mark_with_unique_dimension_labels(m, dt);
m->validate_nodes_and_infer_types();
ov::batch_util::find_batch(m);
if (!track) {
ov::batch_util::restore_original_dimensions(parameter_to_shape, false);
m->validate_nodes_and_infer_types();
return true;
ov::batch_util::restore_original_dimensions(m, parameter_to_shape, false, false);
return false; // we have called validation on this model already
}
ov::batch_util::restore_original_dimensions(parameter_to_shape);
m->validate_nodes_and_infer_types();
ov::batch_util::restore_original_dimensions(m, parameter_to_shape);
bool failed_to_propagate_batch = ov::batch_util::check_batch_tracks_through_all_the_nodes(m);
if (failed_to_propagate_batch) { // restore original input shape with labels
for (const auto& item : parameter_to_shape)
item.first->set_partial_shape(item.second);
} else { // restore original input shape with batch labels
ov::batch_util::restore_original_dimensions(parameter_to_shape, false);
}
m->validate_nodes_and_infer_types();
return true;
ov::batch_util::restore_original_dimensions(m, parameter_to_shape, false, failed_to_propagate_batch);
return false; // we have called validation on this model already
}

View File

@ -93,6 +93,40 @@ TEST(TransformationTests, AutoBatch_FindBatch_Transpose_and_Convolution) {
ASSERT_TRUE(!ov::DimensionTracker::get_label(out_shape[3])) << out_shape;
}
TEST(TransformationTests, AutoBatch_LabelPropagation_Convolution_Reshape) {
auto data = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 4, 6, 8});
const auto& filters = std::make_shared<ov::opset1::Constant>(ov::element::f32, ov::Shape{1, 4, 3, 3});
const auto& conv = std::make_shared<ov::opset1::Convolution>(data,
filters,
ov::Strides{1, 1},
ov::CoordinateDiff{0, 0},
ov::CoordinateDiff{0, 0},
ov::Strides{1, 1});
const auto& reshape =
std::make_shared<ov::opset1::Reshape>(conv,
ov::opset1::Constant::create(ov::element::i64, {3}, {-1, 4, 6}),
false);
const auto& model = std::make_shared<ov::Model>(ov::NodeVector{reshape}, ov::ParameterVector{data});
ov::pass::Manager m;
m.register_pass<ov::pass::InitNodeInfo>();
m.register_pass<ov::pass::FindBatch>();
m.run_passes(model);
ASSERT_NO_THROW(check_rt_info(model));
const auto& shape = data->get_partial_shape();
ASSERT_TRUE(ov::DimensionTracker::get_label(shape[0])) << shape;
ASSERT_TRUE(!ov::DimensionTracker::get_label(shape[1])) << shape;
ASSERT_TRUE(!ov::DimensionTracker::get_label(shape[2])) << shape;
ASSERT_TRUE(!ov::DimensionTracker::get_label(shape[3])) << shape;
const auto& out_shape = model->get_results()[0]->get_output_partial_shape(0);
ASSERT_TRUE(ov::DimensionTracker::get_label(out_shape[0])) << out_shape;
ASSERT_TRUE(!ov::DimensionTracker::get_label(out_shape[1])) << out_shape;
ASSERT_TRUE(!ov::DimensionTracker::get_label(out_shape[2])) << out_shape;
}
TEST(TransformationTests, AutoBatch_FindBatch_SingleMultiply) {
const auto& data = std::make_shared<ov::opset1::Parameter>(ov::element::f32, ov::Shape{1, 4, 10, 10});

View File

@ -31,6 +31,18 @@ TEST(type_prop, static_value_propagation) {
ASSERT_EQ(r->get_shape(), (Shape{1, 2, 3}));
}
TEST(type_prop, reshape_static_dimension_stops_label_propagation_for_auto_batch_case) {
auto shape = ov::PartialShape({1, 1280, 1, 1});
ov::DimensionTracker::set_label(shape[0], 1);
auto param = make_shared<ov::op::v0::Parameter>(element::f32, shape);
auto pattern = op::v0::Constant::create(element::i64, {2}, {-1, 1280});
auto r = make_shared<op::v1::Reshape>(param, pattern, false);
ASSERT_EQ(r->get_element_type(), element::f32);
ASSERT_EQ(r->get_shape(), (Shape{1, 1280}));
ASSERT_EQ(ov::no_label, ov::DimensionTracker::get_label(r->get_output_partial_shape(0)[0]));
}
TEST(type_prop, interval_value_propagation) {
auto param = make_shared<ov::op::v0::Parameter>(element::f32, PartialShape{Dimension(1, 8), 2, 3});
auto shape_of = make_shared<op::v3::ShapeOf>(param);

View File

@ -208,9 +208,7 @@ std::shared_ptr<ov::ICompiledModel> Plugin::compile_model(const std::shared_ptr<
"Auto-batching operates only networks with inputs/outputs batched by 0th dimension");
}
}
const auto& results = cloned_model->get_results();
for (size_t output_id = 0; output_id < results.size(); output_id++) {
const auto& output = results[output_id];
for (const auto& output : cloned_model->get_results()) {
const auto& shape = output->get_output_partial_shape(0);
if (shape.is_dynamic())
OPENVINO_THROW("Auto-batching does not support dynamic networks!");