diff --git a/mindspore/lite/include/registry/node_parser.h b/mindspore/lite/include/registry/node_parser.h index f683fc7701..043a1e8334 100644 --- a/mindspore/lite/include/registry/node_parser.h +++ b/mindspore/lite/include/registry/node_parser.h @@ -22,10 +22,25 @@ #include #include #include "include/registry/parser_context.h" -#include "proto/onnx.pb.h" -#include "proto/caffe.pb.h" -#include "proto/graph.pb.h" -#include "schema/schema_generated.h" + +namespace onnx { +class GraphProto; +class NodeProto; +} // namespace onnx + +namespace caffe { +class LayerParameter; +} // namespace caffe + +namespace tensorflow { +class NodeDef; +} // namespace tensorflow + +namespace tflite { +struct OperatorT; +struct SubGraphT; +struct ModelT; +} // namespace tflite namespace mindspore { namespace ops { diff --git a/mindspore/lite/test/runtest.sh b/mindspore/lite/test/runtest.sh index a56fdb65da..4e1ec60c8c 100644 --- a/mindspore/lite/test/runtest.sh +++ b/mindspore/lite/test/runtest.sh @@ -20,6 +20,7 @@ export LD_LIBRARY_PATH=./:${TENSORRT_PATH}/lib:${CUDA_HOME}/lib64:${LD_LIBRARY_P cp -r ${CUR_DIR}/ut/test_data/* ./ cp -r ${CUR_DIR}/ut/src/runtime/kernel/arm/test_data/* ./ cp -r ${CUR_DIR}/ut/tools/converter/parser/tflite/test_data/* ./ +cp -r ${CUR_DIR}/ut/tools/converter/registry/test_data/* ./ # prepare data for dataset TEST_DATA_DIR=${CUR_DIR}/../../../tests/ut/data/dataset/ cp -fr $TEST_DATA_DIR/testPK ./data @@ -64,6 +65,7 @@ echo 'run common ut tests' ./lite-test --gtest_filter="TestResizeOpenCL*" ./lite-test --gtest_filter="TestSwishOpenCLCI.Fp32CI" ./lite-test --gtest_filter="ModelParserRegistryTest.TestRegistry" +./lite-test --gtest_filter="NodeParserRegistryTest.TestRegistry" ./lite-test --gtest_filter="PassRegistryTest.TestRegistry" ./lite-test --gtest_filter="TestRegistry.TestAdd" ./lite-test --gtest_filter="TestRegistryCustomOp.TestCustomAdd" diff --git a/mindspore/lite/test/ut/tools/converter/registry/model_parser_registry_test.cc b/mindspore/lite/test/ut/tools/converter/registry/model_parser_registry_test.cc index 7da10f5083..0ae2ae02e0 100644 --- a/mindspore/lite/test/ut/tools/converter/registry/model_parser_registry_test.cc +++ b/mindspore/lite/test/ut/tools/converter/registry/model_parser_registry_test.cc @@ -16,7 +16,7 @@ #include #include "common/common_test.h" -#include "ut/tools/converter/registry/model_parser_test.h" +#include "ut/tools/converter/registry/parser/model_parser_test.h" #include "tools/optimizer/common/gllo_utils.h" using mindspore::converter::ConverterParameters; diff --git a/mindspore/lite/test/ut/tools/converter/registry/node_parser_registry_test.cc b/mindspore/lite/test/ut/tools/converter/registry/node_parser_registry_test.cc new file mode 100644 index 0000000000..41010efcb4 --- /dev/null +++ b/mindspore/lite/test/ut/tools/converter/registry/node_parser_registry_test.cc @@ -0,0 +1,79 @@ +/** + * Copyright 2021 Huawei Technologies Co., Ltd + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "api/ir/func_graph.h" +#include "common/common_test.h" +#include "include/registry/model_parser.h" +#include "include/registry/model_parser_registry.h" +#include "include/registry/node_parser_registry.h" +#include "ops/addn.h" +#include "proto/graph.pb.h" + +using mindspore::converter::kFmkTypeTf; +namespace mindspore { +namespace converter { +class AddNodeParser : public NodeParser { + public: + ops::PrimitiveC *Parse(const tensorflow::NodeDef &tf_op, + const std::map &tf_node_map, + std::vector *inputs, int *output_size) override { + auto prim = std::make_unique(); + if (prim == nullptr) { + MS_LOG(ERROR) << "make a shared_ptr failed."; + return nullptr; + } + *output_size = 1; + for (int i = 0; i < tf_op.input_size(); ++i) { + inputs->push_back(tf_op.input(i)); + } + return prim.release(); + } +}; +REG_NODE_PARSER(kFmkTypeTf, Add, std::make_shared()); +} // namespace converter + +class NodeParserRegistryTest : public CommonTest { + public: + NodeParserRegistryTest() = default; + void SetUp() override { + auto model_parser = registry::ModelParserRegistry::GetModelParser(kFmkTypeTf); + if (model_parser == nullptr) { + return; + } + converter::ConverterParameters converter_parameters; + converter_parameters.fmk = kFmkTypeTf; + converter_parameters.model_file = "./tf_add.pb"; + func_graph_ = model_parser->Parse(converter_parameters); + } + api::FuncGraphPtr func_graph_ = nullptr; +}; + +TEST_F(NodeParserRegistryTest, TestRegistry) { + ASSERT_NE(func_graph_, nullptr); + auto node_list = api::FuncGraph::TopoSort(func_graph_->get_return()); + std::vector cnodes; + for (auto &node : node_list) { + if (node->isa()) { + cnodes.push_back(node->cast()); + } + } + ASSERT_EQ(cnodes.size(), 2); + auto cnode = cnodes.front(); + ASSERT_EQ(cnode->size(), 3); + auto prim = GetValueNode>(cnode->input(0)); + ASSERT_NE(prim, nullptr); +} +} // namespace mindspore diff --git a/mindspore/lite/test/ut/tools/converter/registry/model_parser_test.cc b/mindspore/lite/test/ut/tools/converter/registry/parser/model_parser_test.cc similarity index 98% rename from mindspore/lite/test/ut/tools/converter/registry/model_parser_test.cc rename to mindspore/lite/test/ut/tools/converter/registry/parser/model_parser_test.cc index 0566e20b39..8fa8332d0f 100644 --- a/mindspore/lite/test/ut/tools/converter/registry/model_parser_test.cc +++ b/mindspore/lite/test/ut/tools/converter/registry/parser/model_parser_test.cc @@ -14,7 +14,7 @@ * limitations under the License. */ -#include "ut/tools/converter/registry/model_parser_test.h" +#include "ut/tools/converter/registry/parser/model_parser_test.h" #include #include #include "include/errorcode.h" diff --git a/mindspore/lite/test/ut/tools/converter/registry/model_parser_test.h b/mindspore/lite/test/ut/tools/converter/registry/parser/model_parser_test.h similarity index 95% rename from mindspore/lite/test/ut/tools/converter/registry/model_parser_test.h rename to mindspore/lite/test/ut/tools/converter/registry/parser/model_parser_test.h index 8c18659b80..3e3ca02343 100644 --- a/mindspore/lite/test/ut/tools/converter/registry/model_parser_test.h +++ b/mindspore/lite/test/ut/tools/converter/registry/parser/model_parser_test.h @@ -22,7 +22,7 @@ #include #include "include/registry/model_parser.h" #include "include/registry/model_parser_registry.h" -#include "ut/tools/converter/registry/node_parser_test.h" +#include "ut/tools/converter/registry/parser/node_parser_test.h" namespace mindspore { class ModelParserTest : public converter::ModelParser { diff --git a/mindspore/lite/test/ut/tools/converter/registry/node_parser_test.cc b/mindspore/lite/test/ut/tools/converter/registry/parser/node_parser_test.cc similarity index 97% rename from mindspore/lite/test/ut/tools/converter/registry/node_parser_test.cc rename to mindspore/lite/test/ut/tools/converter/registry/parser/node_parser_test.cc index 05ca940eec..63499708a2 100644 --- a/mindspore/lite/test/ut/tools/converter/registry/node_parser_test.cc +++ b/mindspore/lite/test/ut/tools/converter/registry/parser/node_parser_test.cc @@ -14,7 +14,7 @@ * limitations under the License. */ -#include "ut/tools/converter/registry/node_parser_test.h" +#include "ut/tools/converter/registry/parser/node_parser_test.h" #include #include #include diff --git a/mindspore/lite/test/ut/tools/converter/registry/node_parser_test.h b/mindspore/lite/test/ut/tools/converter/registry/parser/node_parser_test.h similarity index 100% rename from mindspore/lite/test/ut/tools/converter/registry/node_parser_test.h rename to mindspore/lite/test/ut/tools/converter/registry/parser/node_parser_test.h diff --git a/mindspore/lite/test/ut/tools/converter/registry/pass_registry_test.cc b/mindspore/lite/test/ut/tools/converter/registry/pass_registry_test.cc index 5c14502680..44a472a6c8 100644 --- a/mindspore/lite/test/ut/tools/converter/registry/pass_registry_test.cc +++ b/mindspore/lite/test/ut/tools/converter/registry/pass_registry_test.cc @@ -26,7 +26,7 @@ #include "ops/custom.h" #include "tools/converter/optimizer_manager.h" #include "tools/optimizer/common/gllo_utils.h" -#include "ut/tools/converter/registry/model_parser_test.h" +#include "ut/tools/converter/registry/parser/model_parser_test.h" using mindspore::converter::ConverterParameters; using mindspore::converter::kFmkTypeCaffe; diff --git a/mindspore/lite/test/ut/tools/converter/registry/test_data/tf_add.pb b/mindspore/lite/test/ut/tools/converter/registry/test_data/tf_add.pb new file mode 100644 index 0000000000..458a4856c7 Binary files /dev/null and b/mindspore/lite/test/ut/tools/converter/registry/test_data/tf_add.pb differ