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 <memory>
18 #include <unordered_map>
19 #include <algorithm>
20 #include "include/api/types.h"
21 #include "include/api/context.h"
22 #include "include/api/dual_abi_helper.h"
23 #include "include/api/callback/callback.h"
24 #include "include/api/metrics/metrics.h"
25 #include "src/litert/lite_model.h"
26 #include "src/litert/inner_context.h"
27 #include "src/litert/inner_allocator.h"
28 #include "src/litert/cxx_api/model/model_impl.h"
29 #include "src/litert/cxx_api/converters.h"
30 #include "src/litert/cxx_api/graph/graph_data.h"
31 #include "src/litert/cxx_api/tensor/tensor_impl.h"
32 #include "src/litert/cxx_api/tensor_utils.h"
33 #include "src/litert/cxx_api/metrics/metrics_adapter.h"
34 #include "src/litert/cxx_api/metrics/metrics_impl.h"
35 #include "src/litert/cxx_api/callback/callback_adapter.h"
36 #include "src/litert/cxx_api/callback/callback_impl.h"
37 #include "src/common/log_adapter.h"
38 #include "src/train/train_session.h"
39 #include "src/train/static_allocator.h"
40
41 namespace mindspore {
CreateTrainSession(std::shared_ptr<Graph::GraphData> graph_data,std::shared_ptr<TrainCfg> cfg,const std::shared_ptr<lite::InnerContext> & context)42 std::shared_ptr<lite::LiteSession> CreateTrainSession(std::shared_ptr<Graph::GraphData> graph_data,
43 std::shared_ptr<TrainCfg> cfg,
44 const std::shared_ptr<lite::InnerContext> &context) {
45 MS_CHECK_TRUE_MSG(graph_data != nullptr, nullptr, "graph data cannot be nullptr");
46 bool is_train_session = graph_data->IsTrainModel();
47 if (is_train_session) {
48 auto model = graph_data->lite_model();
49 if (model == nullptr || model->buf == nullptr) {
50 MS_LOG(ERROR) << "Lite model has been freed.";
51 return nullptr;
52 }
53 std::shared_ptr<lite::LiteSession> shared_session;
54 auto session = new (std::nothrow) lite::TrainSession();
55 if (session == nullptr) {
56 MS_LOG(ERROR) << "create session failed";
57 return nullptr;
58 }
59 shared_session.reset(session);
60
61 context->allocator = std::make_shared<StaticAllocator>();
62 if (context->allocator == nullptr) {
63 MS_LOG(ERROR) << " cannot convert to static allocation";
64 return nullptr;
65 }
66
67 lite::TrainCfg train_cfg;
68 if (cfg != nullptr) {
69 auto status = A2L_ConvertConfig(cfg.get(), &train_cfg);
70 if (status != kSuccess) {
71 MS_LOG(ERROR) << "Failed to convert Config to Lite Config";
72 return nullptr;
73 }
74 }
75
76 auto ret = session->TrainInit(context, &train_cfg);
77 if (ret != mindspore::lite::RET_OK) {
78 MS_LOG(ERROR) << "init session failed";
79 return nullptr;
80 }
81
82 ret = session->CompileTrainGraph(model);
83 if (ret != mindspore::lite::RET_OK) {
84 MS_LOG(WARNING) << "Compiling Train Graph session failed";
85 return nullptr;
86 }
87 return shared_session;
88 }
89 MS_LOG(DEBUG) << "Session is not a train session.";
90 return nullptr;
91 }
92
93 class TrainSupport {
94 public:
TrainSupport()95 TrainSupport() { CreateTrainSessionCallbackHolder(CreateTrainSession); }
~TrainSupport()96 ~TrainSupport() {}
97 };
98
99 TrainSupport support_train_api;
100 } // namespace mindspore
101