From 0772bbd3021cbff515123a7e767c983260101e5e Mon Sep 17 00:00:00 2001 From: Andrew Bakalin Date: Tue, 26 Jan 2021 12:42:36 +0300 Subject: [PATCH] [IE COMMON] Fixes for EliminateUnsqueezeGather transformation (#4013) --- .../eliminate_unsqueeze_gather.cpp | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/inference-engine/src/transformations/src/transformations/common_optimizations/eliminate_unsqueeze_gather.cpp b/inference-engine/src/transformations/src/transformations/common_optimizations/eliminate_unsqueeze_gather.cpp index 3728468339e..3f4e4a5bef6 100644 --- a/inference-engine/src/transformations/src/transformations/common_optimizations/eliminate_unsqueeze_gather.cpp +++ b/inference-engine/src/transformations/src/transformations/common_optimizations/eliminate_unsqueeze_gather.cpp @@ -13,8 +13,10 @@ NGRAPH_RTTI_DEFINITION(ngraph::pass::EliminateUnsqueezeGather, "EliminateUnsquee ngraph::pass::EliminateUnsqueezeGather::EliminateUnsqueezeGather() { MATCHER_SCOPE(EliminateUnsqueezeGather); + // Remove Unsqueeze + Gather pair, if Gather gathers data by `1` dimension that was previously added by Unsqueeze const auto unsqueezeAxis = ngraph::pattern::any_input(); - const auto unsqueeze = ngraph::pattern::wrap_type({ngraph::pattern::any_input(), unsqueezeAxis}); + const auto unsqueezeInput = ngraph::pattern::any_input(); + const auto unsqueeze = ngraph::pattern::wrap_type({unsqueezeInput, unsqueezeAxis}, pattern::consumers_count(1)); const auto gatherIndices = ngraph::opset6::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0}); const auto gatherAxis = ngraph::pattern::any_input(); const auto gather = ngraph::pattern::wrap_type({unsqueeze, gatherIndices, gatherAxis}); @@ -45,14 +47,14 @@ ngraph::pass::EliminateUnsqueezeGather::EliminateUnsqueezeGather() { auto& m_gather = patternValue.at(gather); const auto& m_unsqueeze = patternValue.at(unsqueeze); - const auto& unsqueezeData = m_unsqueeze.get_node_shared_ptr()->get_input_node_shared_ptr(0); + const auto& m_unsqueezeInput = patternValue.at(unsqueezeInput); - ngraph::copy_runtime_info(m_gather.get_node_shared_ptr(), unsqueezeData); - m_gather.replace(unsqueezeData); + ngraph::copy_runtime_info(m_gather.get_node_shared_ptr(), m_unsqueeze.get_node_shared_ptr()); + m_gather.replace(m_unsqueezeInput); return true; }; auto m = std::make_shared(gather, "EliminateUnsqueezeGather"); - ngraph::pass::MatcherPass::register_matcher(m, callback); + register_matcher(m, callback); }