diff --git a/include/api/serialization.h b/include/api/serialization.h index dcb0a4762ae..613355df152 100644 --- a/include/api/serialization.h +++ b/include/api/serialization.h @@ -69,7 +69,8 @@ class MS_API Serialization { static Status SetParameters(const std::map ¶meters, Model *model); static Status ExportModel(const Model &model, ModelType model_type, Buffer *model_data); static Status ExportModel(const Model &model, ModelType model_type, const std::string &model_file, - QuantizationType quantization_type = kNoQuant, bool export_inference_only = true); + QuantizationType quantization_type = kNoQuant, bool export_inference_only = true, + std::vector output_tensor_name = {}); private: static Status Load(const void *model_data, size_t data_size, ModelType model_type, Graph *graph, const Key &dec_key, diff --git a/mindspore/ccsrc/cxx_api/serialization.cc b/mindspore/ccsrc/cxx_api/serialization.cc index e14dab3a9a4..ec45b813f69 100644 --- a/mindspore/ccsrc/cxx_api/serialization.cc +++ b/mindspore/ccsrc/cxx_api/serialization.cc @@ -278,7 +278,8 @@ Status Serialization::ExportModel(const Model &, ModelType, Buffer *) { return kMEFailed; } -Status Serialization::ExportModel(const Model &, ModelType, const std::string &, QuantizationType, bool) { +Status Serialization::ExportModel(const Model &, ModelType, const std::string &, QuantizationType, bool, + std::vector output_tensor_name) { MS_LOG(ERROR) << "Unsupported feature."; return kMEFailed; } diff --git a/mindspore/lite/include/lite_session.h b/mindspore/lite/include/lite_session.h index 6827a1315ea..9e6edfc3bcb 100644 --- a/mindspore/lite/include/lite_session.h +++ b/mindspore/lite/include/lite_session.h @@ -182,9 +182,11 @@ class MS_API LiteSession { /// \param[in] model_type indication whether to save full model or only the inference part /// \param[in] quant_type indication whether to quantize exported model /// \param[in] format of exported file (currently only FT_FLATBUFFERS is supported) + /// \param[in] out_put_tensor_name of exported tensorname /// \return STATUS as an error code of the set operation, STATUS is defined in errorcode.h virtual int Export(const std::string &file_name, lite::ModelType model_type = lite::MT_TRAIN, - lite::QuantizationType quant_type = lite::QT_DEFAULT, lite::FormatType = lite::FT_FLATBUFFERS) { + lite::QuantizationType quant_type = lite::QT_DEFAULT, lite::FormatType = lite::FT_FLATBUFFERS, + std::vector out_put_tensor_name = {}) { return mindspore::lite::RET_ERROR; } diff --git a/mindspore/lite/src/cxx_api/serialization.cc b/mindspore/lite/src/cxx_api/serialization.cc index d601c713af9..a7c1796fea4 100644 --- a/mindspore/lite/src/cxx_api/serialization.cc +++ b/mindspore/lite/src/cxx_api/serialization.cc @@ -123,7 +123,8 @@ Status Serialization::ExportModel(const Model &model, ModelType model_type, Buff } Status Serialization::ExportModel(const Model &model, ModelType model_type, const std::string &model_file, - QuantizationType quantization_type, bool export_inference_only) { + QuantizationType quantization_type, bool export_inference_only, + std::vector output_tensor_name) { if (model.impl_ == nullptr) { MS_LOG(ERROR) << "Model implement is null."; return kLiteUninitializedObj; @@ -137,7 +138,7 @@ Status Serialization::ExportModel(const Model &model, ModelType model_type, cons return kLiteParamInvalid; } auto ret = model.impl_->session_->Export(model_file, export_inference_only ? lite::MT_INFERENCE : lite::MT_TRAIN, - A2L_ConvertQT(quantization_type), lite::FT_FLATBUFFERS); + A2L_ConvertQT(quantization_type), lite::FT_FLATBUFFERS, output_tensor_name); return (ret == mindspore::lite::RET_OK) ? kSuccess : kLiteError; } diff --git a/mindspore/lite/src/train/train_export.cc b/mindspore/lite/src/train/train_export.cc index 6872ed0a1bb..b91d08b00c6 100644 --- a/mindspore/lite/src/train/train_export.cc +++ b/mindspore/lite/src/train/train_export.cc @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include #include #include @@ -28,8 +30,42 @@ namespace mindspore { namespace lite { +namespace { constexpr static int kFmkVal = 3; constexpr static int kTransformTensorDim = 4; +std::vector GetLinkedPostIdx(const schema::MetaGraphT &graphT, const size_t &tensorIdx) { + std::vector postNodeIdx; + for (size_t i = 0; i < graphT.nodes.size(); i++) { + auto &oldNode = graphT.nodes.at(i); + if (oldNode == nullptr) { + continue; + } + auto inputIndexes = oldNode->inputIndex; + if (IsContain(inputIndexes, tensorIdx)) { + postNodeIdx.emplace_back(i); + } + } + return postNodeIdx; +} + +std::vector GetOutputNodeIdx(const schema::MetaGraphT &graphT, const schema::CNodeT &node, + const int outputIndexIdx = -1) { + std::vector outputIndexes; + if (outputIndexIdx == -1) { + outputIndexes = node.outputIndex; + } else { + outputIndexes.emplace_back(node.outputIndex.at(outputIndexIdx)); + } + std::set outputNodeIdx; + for (uint32_t outputIdx : outputIndexes) { + auto linkedPostIdx = GetLinkedPostIdx(graphT, outputIdx); + outputNodeIdx.insert(linkedPostIdx.begin(), linkedPostIdx.end()); + } + std::vector ret; + ret.insert(ret.end(), outputNodeIdx.begin(), outputNodeIdx.end()); + return ret; +} +} // namespace std::vector TrainExport::CreateData(const lite::Tensor *tensor) { uint8_t *tensor_data = reinterpret_cast(tensor->data_c()); @@ -397,16 +433,90 @@ int TrainExport::ExportNet(const std::vector &k } } TagQuantizedNodes(); // do another loop to mark QUANT_WEIGHT_NODES + auto status = TopologicalSort(); + if (status != RET_OK) { + MS_LOG(ERROR) << "TopologicalSort failed."; + return RET_ERROR; + } return RET_OK; } +int TrainExport::TopologicalSort() { + MS_ASSERT(meta_graph_ != nullptr); + std::vector> new_nodes; + std::vector sinked_tensor_idxes; + for (auto &subgraph : meta_graph_->subGraph) { + std::copy(subgraph->inputIndices.begin(), subgraph->inputIndices.end(), std::back_inserter(sinked_tensor_idxes)); + } + // put all const tensor index into sinked_tensor_idxes + for (size_t i = 0; i < meta_graph_->allTensors.size(); i++) { + if (meta_graph_->allTensors.at(i)->nodeType == NodeType_ValueNode) { + sinked_tensor_idxes.push_back(i); + } + } + auto &old_nodes = meta_graph_->nodes; + std::queue> op_queue; + // put all none depend node into queue + for (size_t i = 0; i < meta_graph_->subGraph.size(); i++) { + std::vector new_subgraph_node_indices = {}; + auto subgraph_node_indices = meta_graph_->subGraph[i]->nodeIndices; + + for (size_t j = 0; j < subgraph_node_indices.size(); j++) { + auto &node = old_nodes[subgraph_node_indices[j]]; + if (IsNodeNonDepend(node, sinked_tensor_idxes)) { + sinked_tensor_idxes.insert(sinked_tensor_idxes.end(), node->outputIndex.begin(), node->outputIndex.end()); + op_queue.push(std::move(node)); + } + } + while (!op_queue.empty()) { + auto &node = op_queue.front(); + auto post_node_idxes = GetOutputNodeIdx(*meta_graph_, *(node.get())); + sinked_tensor_idxes.insert(sinked_tensor_idxes.end(), node->outputIndex.begin(), node->outputIndex.end()); + for (auto post_node_idx : post_node_idxes) { + if (IsContain(subgraph_node_indices, (unsigned int)(post_node_idx))) { + auto &post_node = old_nodes.at(post_node_idx); + // check if post_node is non-depended + if (IsNodeNonDepend(post_node, sinked_tensor_idxes)) { + op_queue.push(std::move(post_node)); + } + } + } + new_nodes.emplace_back(std::move(node)); + new_subgraph_node_indices.push_back(new_nodes.size() - 1); + op_queue.pop(); + } + meta_graph_->subGraph[i]->nodeIndices.swap(new_subgraph_node_indices); + } + if (new_nodes.size() != old_nodes.size()) { + MS_LOG(ERROR) << "Unknown error in TopologicalSort, old_nodes size: " << old_nodes.size() + << ", new_nodes size: " << new_nodes.size(); + return RET_ERROR; + } + meta_graph_->nodes.swap(new_nodes); + return RET_OK; +} + +bool TrainExport::IsNodeNonDepend(const std::unique_ptr &node, + const std::vector &sinked_tensor_idxes) { + MS_ASSERT(node != nullptr); + return std::all_of(node->inputIndex.begin(), node->inputIndex.end(), + [&](size_t input_idx) { return IsContain(sinked_tensor_idxes, size_t(input_idx)); }); +} + int TrainExport::ExportInit(const std::string model_name, std::string version) { meta_graph_ = new (std::nothrow) schema::MetaGraphT(); if (meta_graph_ == nullptr) { MS_LOG(ERROR) << "cannot allocate meta_graph"; return RET_ERROR; } + auto sub_graph = std::make_unique(); + if (sub_graph == nullptr) { + MS_LOG(ERROR) << "cannot allocate SubGraphT"; + return RET_ERROR; + } + sub_graph->name = model_name + "_subgraph"; + meta_graph_->subGraph.emplace_back(std::move(sub_graph)); meta_graph_->fmkType = kFmkVal; meta_graph_->name = model_name; meta_graph_->version = version; diff --git a/mindspore/lite/src/train/train_export.h b/mindspore/lite/src/train/train_export.h index fc31c13b467..477d71a5d5f 100644 --- a/mindspore/lite/src/train/train_export.h +++ b/mindspore/lite/src/train/train_export.h @@ -57,6 +57,8 @@ class TrainExport { std::vector out_idx_; std::map remap_; std::unordered_map connect_; // connection map (backbone tenor id-> head tensor id) + bool IsNodeNonDepend(const std::unique_ptr &node, const std::vector &sinked_tensor_idxes); + int TopologicalSort(); void PrepareRemap(int offset); Model::Node *FindNode(const mindspore::kernel::LiteKernel *kernel, const Model *model); std::unique_ptr CreateTensor(const Tensor *tensor, schema::Tensor *scTensor); diff --git a/mindspore/lite/src/train/train_session.cc b/mindspore/lite/src/train/train_session.cc index 8a47ca6ab4f..079a6657a3a 100644 --- a/mindspore/lite/src/train/train_session.cc +++ b/mindspore/lite/src/train/train_session.cc @@ -22,6 +22,7 @@ #include #include #include +#include #include #include "include/errorcode.h" #include "src/executor.h" @@ -793,8 +794,63 @@ int TrainSession::Resize(const std::vector &inputs, const st return RET_OK; } +int TrainSession::FindUseInTensorKernel(std::vector *use_in_tensor_kernels, + const std::vector &kernel_in_tensors, + const std::vector &inference_kernels) { + for (size_t i = 0; i < inference_kernels.size(); i++) { + for (size_t j = 0; j < kernel_in_tensors.size(); j++) { + if (IsContain(inference_kernels[i]->out_tensors(), kernel_in_tensors[j])) { + use_in_tensor_kernels->push_back(inference_kernels[i]); + } + } + } + return RET_OK; +} + +int TrainSession::FindExportKernels(std::vector *export_kernels, + const std::vector &export_output_tensor_names, + const std::vector &inference_kernels) { + std::vector all_kernel_name = {}; + std::transform(inference_kernels.begin(), inference_kernels.end(), std::back_inserter(all_kernel_name), + [](kernel::LiteKernel *kernel) { return kernel->name(); }); + std::queue need_kernel_names; + // Find the kernel name according to the tensor name + for (auto &kernel : inference_kernels) { + if (std::any_of(kernel->out_tensors().begin(), kernel->out_tensors().end(), [&](lite::Tensor *out_tensor) { + return IsContain(export_output_tensor_names, out_tensor->tensor_name()); + })) { + need_kernel_names.push(kernel->name()); + } + } + // find all kernel + while (!need_kernel_names.empty()) { + auto kernel_name = need_kernel_names.front(); + need_kernel_names.pop(); + auto it = find(all_kernel_name.begin(), all_kernel_name.end(), kernel_name); + if (it == all_kernel_name.end()) { + MS_LOG(ERROR) << "not find kernel name in export trained model."; + return RET_ERROR; + } + auto kernel = inference_kernels[it - all_kernel_name.begin()]; + if (!IsContain(*export_kernels, kernel)) { + export_kernels->push_back(kernel); + } + auto kernel_in_tensors = kernel->in_tensors(); + std::vector use_in_tensor_kernels; + auto status = FindUseInTensorKernel(&use_in_tensor_kernels, kernel_in_tensors, inference_kernels); + if (status != RET_OK) { + MS_LOG(ERROR) << "FindUseInTensorKernel failed."; + return RET_ERROR; + } + for (size_t i = 0; i < use_in_tensor_kernels.size(); i++) { + need_kernel_names.push(use_in_tensor_kernels[i]->name()); + } + } + return RET_OK; +} + int TrainSession::Export(const std::string &file_name, ModelType model_type, QuantizationType quant_type, - FormatType format) { + FormatType format, std::vector out_put_tensor_name) { if (file_name.empty()) { MS_LOG(ERROR) << "File name cannot be empty"; return RET_ERROR; @@ -820,9 +876,21 @@ int TrainSession::Export(const std::string &file_name, ModelType model_type, Qua MS_LOG(ERROR) << "cannot init export"; return status; } - status = texport.ExportNet((model_type == MT_TRAIN) ? train_kernels_ : inference_kernels_, tensors_, - (model_type == MT_TRAIN) ? train_output_tensor_names_ : eval_output_tensor_names_, - model_.get(), quant_type); + + if (!out_put_tensor_name.empty() && model_type == MT_INFERENCE) { + std::vector export_kernels = {}; + status = FindExportKernels(&export_kernels, out_put_tensor_name, inference_kernels_); + if (status != RET_OK) { + MS_LOG(ERROR) << "FindExportKernels failed."; + return RET_ERROR; + } + status = texport.ExportNet(export_kernels, tensors_, out_put_tensor_name, model_.get(), quant_type); + } else { + status = texport.ExportNet((model_type == MT_TRAIN) ? train_kernels_ : inference_kernels_, tensors_, + (model_type == MT_TRAIN) ? train_output_tensor_names_ : eval_output_tensor_names_, + model_.get(), quant_type); + } + if (status != RET_OK) { MS_LOG(ERROR) << "cannot export Network"; return status; diff --git a/mindspore/lite/src/train/train_session.h b/mindspore/lite/src/train/train_session.h index ee7f2863ef5..5878add29d1 100644 --- a/mindspore/lite/src/train/train_session.h +++ b/mindspore/lite/src/train/train_session.h @@ -89,11 +89,18 @@ class TrainSession : virtual public lite::LiteSession { } return outputs; } - int Export(const std::string &fb_name, ModelType model_type, QuantizationType quant_type, FormatType) override; + int Export(const std::string &fb_name, ModelType model_type, QuantizationType quant_type, FormatType, + std::vector out_put_tensor_name = {}) override; std::vector GetFeatureMaps() const override; int UpdateFeatureMaps(const std::vector &features_map) override; + int FindUseInTensorKernel(std::vector *use_in_tensor_kernels, + const std::vector &kernel_in_tensors, + const std::vector &inference_kernels); + int FindExportKernels(std::vector *export_kernels, + const std::vector &export_output_tensor_names, + const std::vector &inference_kernels); protected: int AllocWorkSpace(); diff --git a/mindspore/lite/src/train/transfer_session.cc b/mindspore/lite/src/train/transfer_session.cc index 4234d7a6fae..28532cfeb43 100644 --- a/mindspore/lite/src/train/transfer_session.cc +++ b/mindspore/lite/src/train/transfer_session.cc @@ -179,7 +179,7 @@ std::unordered_map TransferSession::ConnectionMap() { } int TransferSession::Export(const std::string &filename, ModelType model_type, QuantizationType quant_type, - FormatType format) { + FormatType format, std::vector out_put_tensor_name) { if (format != FT_FLATBUFFERS) { MS_LOG(ERROR) << "Currently only flatbuffer format is supported"; return RET_ERROR; @@ -206,7 +206,17 @@ int TransferSession::Export(const std::string &filename, ModelType model_type, Q return status; } } - status = texport.ExportNet(inference_kernels_, tensors_, GetOutputTensorNames(), model_.get(), quant_type); + if (!out_put_tensor_name.empty() && model_type == MT_INFERENCE) { + std::vector export_kernels = {}; + status = FindExportKernels(&export_kernels, out_put_tensor_name, inference_kernels_); + if (status != RET_OK) { + MS_LOG(ERROR) << "FindExportKernels failed."; + return RET_ERROR; + } + status = texport.ExportNet(export_kernels, tensors_, out_put_tensor_name, model_.get(), quant_type); + } else { + status = texport.ExportNet(inference_kernels_, tensors_, GetOutputTensorNames(), model_.get(), quant_type); + } if (status != RET_OK) { MS_LOG(ERROR) << "cannot serialize head"; return status; diff --git a/mindspore/lite/src/train/transfer_session.h b/mindspore/lite/src/train/transfer_session.h index 3d76a205a94..04fca0e1582 100644 --- a/mindspore/lite/src/train/transfer_session.h +++ b/mindspore/lite/src/train/transfer_session.h @@ -61,7 +61,8 @@ class TransferSession : public lite::TrainSession { mindspore::tensor::MSTensor *GetInputsByTensorName(const std::string &tensor_name) const override; int CompileTransferGraph(); - int Export(const std::string &fb_name, ModelType model_type, QuantizationType quant_type, FormatType) override; + int Export(const std::string &fb_name, ModelType model_type, QuantizationType quant_type, FormatType, + std::vector out_put_tensor_name = {}) override; protected: lite::LiteSession *backbone_session_ = nullptr;