[VPU][NGraph] Extend TopK K propagation (#2943)

* [VPU][NGraph] Support ShapeOf and Gather in TopK K propagation

* [VPU] Save calculated K value

* [VPU][Tests] Introduces tests

* [Tests] Review fixes
This commit is contained in:
Andrew Bakalin
2020-11-13 07:50:31 +03:00
committed by GitHub
parent 69a8195d56
commit 0681cdf1c5
4 changed files with 162 additions and 48 deletions

View File

@@ -58,6 +58,7 @@ public:
protected:
int64_t m_axis;
int64_t m_maximumK;
uint64_t m_normalized_axis;
Mode m_mode;
SortType m_sort;

View File

@@ -18,6 +18,7 @@ ngraph::vpu::op::StaticShapeTopK::StaticShapeTopK(
const element::Type& index_element_type)
: Op{{data, k}}
, m_axis{axis}
, m_maximumK{-1}
, m_normalized_axis{0}
, m_mode{as_enum<Mode>(mode)}
, m_sort{as_enum<SortType>(sort)}
@@ -34,6 +35,7 @@ ngraph::vpu::op::StaticShapeTopK::StaticShapeTopK(
const ngraph::element::Type &index_element_type)
: Op{{data, k}}
, m_axis{axis}
, m_maximumK{-1}
, m_normalized_axis{0}
, m_mode{mode}
, m_sort{sort}
@@ -95,11 +97,12 @@ void ngraph::vpu::op::StaticShapeTopK::validate_and_infer_types() {
const auto is_max_value_calculated = max_k.first;
const auto calculated_max_value = max_k.second;
if (is_max_value_calculated) {
output_shape[m_normalized_axis] = calculated_max_value;
} else {
output_shape[m_normalized_axis] = -1;
m_maximumK = calculated_max_value;
}
}
output_shape[m_normalized_axis] = m_maximumK;
NODE_VALIDATION_CHECK(this, output_shape.is_static(),
"StaticShapeTopK output shape is not fully defined: ", output_shape);

View File

@@ -7,7 +7,7 @@
#include <ngraph/type/element_type.hpp>
#include <ngraph/function.hpp>
#include <ngraph_functions/utils/ngraph_helpers.hpp>
#include <ngraph/opsets/opset4.hpp>
#include <ngraph/opsets/opset5.hpp>
#include <vpu/ngraph/operations/dynamic_shape_resolver.hpp>
#include <vpu/ngraph/operations/static_shape_topk.hpp>
#include <vpu/ngraph/transformations/dynamic_to_static_shape_topk.hpp>
@@ -15,34 +15,14 @@
namespace {
class DynamicToStaticTopKPropagation : public CommonTestUtils::TestsCommon,
public testing::WithParamInterface<int64_t> {
public:
void SetUp() override {
class DynamicToStaticTopKPropagationBase : public CommonTestUtils::TestsCommon,
public testing::WithParamInterface<int64_t> {
protected:
void validate(const ngraph::Function& function) const {
const auto& k = GetParam();
const auto data = std::make_shared<ngraph::opset4::Parameter>(ngraph::element::i64, ngraph::Shape{1000});
const auto realK = std::make_shared<ngraph::opset4::Parameter>(ngraph::element::i32, ngraph::Shape{1});
const auto maxK = ngraph::opset4::Constant::create(ngraph::element::i32, {1}, {k});
const auto concat = std::make_shared<ngraph::opset4::Concat>(ngraph::OutputVector{realK, maxK}, 0);
const auto reduceMin = std::make_shared<ngraph::opset4::ReduceMin>(concat, ngraph::opset4::Constant::create(ngraph::element::i32, {1}, {0}), false);
const auto builtSubgraph = buildSubgraph(reduceMin);
const auto dsr = std::make_shared<ngraph::vpu::op::DynamicShapeResolver>(data, realK);
const auto topK = std::make_shared<ngraph::opset4::TopK>(dsr, builtSubgraph, 0, "max", "value");
ngraph::ResultVector results{std::make_shared<ngraph::opset4::Result>(topK->output(0)),
std::make_shared<ngraph::opset4::Result>(topK->output(1))};
const auto function = std::make_shared<ngraph::Function>(results, ngraph::ParameterVector{data, realK}, "TopKPropagationOfK");
topK->set_output_type(0, dsr->get_input_element_type(0), ngraph::PartialShape::dynamic(1));
const auto transformations = vpu::Transformations{{topK->type_info, vpu::dynamicToStaticShapeTopK}};
ASSERT_NO_THROW(vpu::DynamicToStaticShape(transformations).run_on_function(function));
ngraph::ResultVector processedResults;
ASSERT_NO_THROW(processedResults = function->get_results());
ASSERT_NO_THROW(processedResults = function.get_results());
EXPECT_EQ(processedResults.size(), 2);
const auto topKOutPartialShape = processedResults[0]->get_input_partial_shape(0);
@@ -53,55 +33,133 @@ public:
EXPECT_EQ(topKOutShape[0], k);
}
protected:
virtual std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const {
return node;
}
};
class DynamicToStaticTopKPropagationConcatBased : public DynamicToStaticTopKPropagationBase {
public:
void SetUp() override {
const auto& k = GetParam();
ngraph::Shape inputShape{upperBoundK};
const auto data = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, inputShape);
const auto dataShape = ngraph::opset5::Constant::create(ngraph::element::i32, ngraph::Shape{inputShape.size()}, inputShape);
const auto resultK = ngraph::opset5::Constant::create(ngraph::element::i32, {1}, {k});
const auto concat = std::make_shared<ngraph::opset5::Concat>(ngraph::OutputVector{dataShape, resultK}, 0);
const auto reduceMin = std::make_shared<ngraph::opset5::ReduceMin>(concat, ngraph::opset5::Constant::create(ngraph::element::i32, {1}, {0}), false);
const auto builtSubgraph = buildSubgraph(reduceMin);
const auto dsr = std::make_shared<ngraph::vpu::op::DynamicShapeResolver>(data, dataShape);
const auto topK = std::make_shared<ngraph::opset5::TopK>(dsr, builtSubgraph, 0, "max", "value");
ngraph::ResultVector results{std::make_shared<ngraph::opset5::Result>(topK->output(0)),
std::make_shared<ngraph::opset5::Result>(topK->output(1))};
const auto function = std::make_shared<ngraph::Function>(results, ngraph::ParameterVector{data}, "TopKPropagationOfK");
const auto transformations = vpu::Transformations{{topK->type_info, vpu::dynamicToStaticShapeTopK}};
ASSERT_NO_THROW(vpu::DynamicToStaticShape(transformations).run_on_function(function));
validate(*function);
}
protected:
static constexpr int64_t upperBoundK = 1000;
};
const std::vector<int64_t> kVec = {0, 10, 100, 200, 500};
TEST_P(DynamicToStaticTopKPropagation, KPropagation) {
TEST_P(DynamicToStaticTopKPropagationConcatBased, KPropagation) {
}
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagation, ::testing::ValuesIn(kVec));
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatBased, ::testing::ValuesIn(kVec));
class DynamicToStaticTopKPropagationReshape : public DynamicToStaticTopKPropagation {
class DynamicToStaticTopKPropagationConcatReshape : public DynamicToStaticTopKPropagationConcatBased {
protected:
std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
return std::make_shared<ngraph::opset4::Reshape>(node, ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}), false);
return std::make_shared<ngraph::opset5::Reshape>(node, ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}), false);
}
};
TEST_P(DynamicToStaticTopKPropagationReshape, KPropagation) {
TEST_P(DynamicToStaticTopKPropagationConcatReshape, KPropagation) {
}
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationReshape, ::testing::ValuesIn(kVec));
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatReshape, ::testing::ValuesIn(kVec));
class DynamicToStaticTopKPropagationSqueezeUnsqueeze : public DynamicToStaticTopKPropagation {
class DynamicToStaticTopKPropagationConcatSqueezeUnsqueeze : public DynamicToStaticTopKPropagationConcatBased {
protected:
std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
const auto unsqueeze = std::make_shared<ngraph::opset4::Unsqueeze>(node, ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
return std::make_shared<ngraph::opset4::Squeeze>(unsqueeze, ngraph::opset4::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
const auto unsqueeze = std::make_shared<ngraph::opset5::Unsqueeze>(
node,
ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
return std::make_shared<ngraph::opset5::Squeeze>(unsqueeze, ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{1}, {0}));
}
};
TEST_P(DynamicToStaticTopKPropagationSqueezeUnsqueeze, KPropagation) {
TEST_P(DynamicToStaticTopKPropagationConcatSqueezeUnsqueeze, KPropagation) {
}
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationSqueezeUnsqueeze, ::testing::ValuesIn(kVec));
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatSqueezeUnsqueeze, ::testing::ValuesIn(kVec));
class DynamicToStaticTopKPropagationConvert : public DynamicToStaticTopKPropagation {
class DynamicToStaticTopKPropagationConcatConvert : public DynamicToStaticTopKPropagationConcatBased {
protected:
std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
const auto convert = std::make_shared<ngraph::opset4::Convert>(node, ngraph::element::i32);
return std::make_shared<ngraph::opset4::Convert>(convert, ngraph::element::i64);
const auto convert = std::make_shared<ngraph::opset5::Convert>(node, ngraph::element::i32);
return std::make_shared<ngraph::opset5::Convert>(convert, ngraph::element::i64);
}
};
TEST_P(DynamicToStaticTopKPropagationConvert, KPropagation) {
TEST_P(DynamicToStaticTopKPropagationConcatConvert, KPropagation) {
}
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConvert, ::testing::ValuesIn(kVec));
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationConcatConvert, ::testing::ValuesIn(kVec));
class DynamicToStaticTopKPropagationShapeOfBased : public DynamicToStaticTopKPropagationBase {
public:
void SetUp() override {
const auto& k = GetParam();
const auto upperBoundData = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i64, ngraph::Shape{upperBoundK});
const auto realData = std::make_shared<ngraph::opset5::Parameter>(ngraph::element::i32, ngraph::Shape{static_cast<size_t>(k)});
const auto dsr = std::make_shared<ngraph::vpu::op::DynamicShapeResolver>(
upperBoundData,
ngraph::opset5::Constant::create(ngraph::element::i32, ngraph::Shape{1}, {k}));
const auto shapeOf = std::make_shared<ngraph::opset5::ShapeOf>(realData);
const auto builtSubgraph = buildSubgraph(shapeOf);
const auto topK = std::make_shared<ngraph::opset5::TopK>(dsr, builtSubgraph, 0, "max", "value");
ngraph::ResultVector results{std::make_shared<ngraph::opset5::Result>(topK->output(0)),
std::make_shared<ngraph::opset5::Result>(topK->output(1))};
const auto function = std::make_shared<ngraph::Function>(results, ngraph::ParameterVector{upperBoundData, realData}, "TopKPropagationOfK");
const auto transformations = vpu::Transformations{{topK->type_info, vpu::dynamicToStaticShapeTopK}};
ASSERT_NO_THROW(vpu::DynamicToStaticShape(transformations).run_on_function(function));
validate(*function);
}
protected:
static constexpr int64_t upperBoundK = 1000;
};
class DynamicToStaticTopKPropagationShapeOfGather : public DynamicToStaticTopKPropagationShapeOfBased {
protected:
std::shared_ptr<ngraph::Node> buildSubgraph(std::shared_ptr<ngraph::Node> node) const override {
return std::make_shared<ngraph::opset5::Gather>(
node,
ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0}),
ngraph::opset5::Constant::create(ngraph::element::i64, ngraph::Shape{}, {0}));
}
};
TEST_P(DynamicToStaticTopKPropagationShapeOfGather, KPropagation) {
}
INSTANTIATE_TEST_CASE_P(smoke_NGraph, DynamicToStaticTopKPropagationShapeOfGather, ::testing::ValuesIn(kVec));
} // namespace