• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2021 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include "include/registry/model_parser_registry.h"
18 #include <map>
19 #include "src/common/log_adapter.h"
20 #include "nnacl/op_base.h"
21 
22 namespace mindspore {
23 namespace registry {
24 namespace {
25 std::map<FmkType, ModelParserCreator> model_parser_room;
26 }  // namespace
27 
ModelParserRegistry(FmkType fmk,ModelParserCreator creator)28 ModelParserRegistry::ModelParserRegistry(FmkType fmk, ModelParserCreator creator) {
29   if ((fmk < converter::kFmkTypeTf || fmk >= converter::kFmkTypeEnd) && fmk != converter::kFmkTypeOM) {
30     MS_LOG(ERROR) << "FMK must be in FmkType!Illegal FMK: " << fmk;
31     return;
32   }
33   if (creator == nullptr) {
34     MS_LOG(ERROR) << "ILLEGAL CREATOR: model parser creator is nullptr.";
35     return;
36   }
37   model_parser_room[fmk] = creator;
38 }
39 
GetModelParser(FmkType fmk)40 converter::ModelParser *ModelParserRegistry::GetModelParser(FmkType fmk) {
41   if ((fmk < converter::kFmkTypeTf || fmk >= converter::kFmkTypeEnd) && fmk != converter::kFmkTypeOM) {
42     MS_LOG(ERROR) << "ILLEGAL FMK: fmk must be in FmkType.";
43     return nullptr;
44   }
45   auto it = model_parser_room.find(fmk);
46   if (it != model_parser_room.end()) {
47     auto creator = it->second;
48     MS_CHECK_TRUE_RET(creator != nullptr, nullptr);
49     return creator();
50   }
51   return nullptr;
52 }
53 }  // namespace registry
54 }  // namespace mindspore
55