1 /** 2 * Copyright 2020 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 #ifndef MINDSPORE_LITE_SRC_TRAIN_TRANSFER_SESSION_H_ 17 #define MINDSPORE_LITE_SRC_TRAIN_TRANSFER_SESSION_H_ 18 #include <vector> 19 #include <string> 20 #include <tuple> 21 #include <unordered_map> 22 #include <utility> 23 #include "src/lite_session.h" 24 #include "src/train/train_session.h" 25 26 /* 27 Inheritance Diagram 28 29 +-------------------------------+ 30 | session::LiteSession | 31 +--------------↑----------------+ 32 | 33 +--------------+----------------+ 34 | lite::LiteSession | 35 +--------------↑----------------+ 36 | 37 +--------------+----------------+ 38 | lite::TrainSession | 39 +--------------↑----------------+ 40 | 41 +--------------+----------------+ 42 | lite::TrasferSession | 43 +-------------------------------+ 44 */ 45 46 namespace mindspore { 47 namespace lite { 48 49 class TransferSession : public lite::TrainSession { 50 public: 51 explicit TransferSession(const char *model_buf_backbone, size_t size_backbone, const lite::Context *context); 52 53 ~TransferSession(); 54 is_valid()55 bool is_valid() const { return is_valid_; } 56 57 int RunGraph(const KernelCallBack &before = nullptr, const KernelCallBack &after = nullptr) override; 58 59 void BindThread(bool if_bind) override; 60 std::vector<tensor::MSTensor *> GetInputs() const override; 61 mindspore::tensor::MSTensor *GetInputsByTensorName(const std::string &tensor_name) const override; 62 63 int CompileTransferGraph(); 64 int Export(const std::string &fb_name, ModelType model_type, QuantizationType quant_type, FormatType, 65 std::vector<std::string> out_put_tensor_name = {}) override; 66 67 protected: 68 lite::LiteSession *backbone_session_ = nullptr; 69 char *lite_model_ = nullptr; 70 std::vector<mindspore::tensor::MSTensor *> combined_inputs_; 71 std::vector<std::pair<mindspore::tensor::MSTensor *, mindspore::tensor::MSTensor *>> backbone_head_map_; 72 bool is_valid_ = false; 73 74 private: 75 bool CompileFormatTransform(tensor::MSTensor *out, tensor::MSTensor *in, int *mask, size_t mask_len); 76 std::unordered_map<size_t, size_t> ConnectionMap(); 77 bool nchw2nhwc_ = false; 78 size_t size_backbone_; 79 }; 80 } // namespace lite 81 } // namespace mindspore 82 #endif // MINDSPORE_LITE_SRC_TRAIN_TRANSFER_SESSION_H_ 83