yankf21
/
mindspore2022
forked from huawei/mindspore2022
0
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

!19186 [MS][LITE][TOD] Fixed exporting of albert model 

Merge pull request !19186 from ehaleva/i3yl2h
This commit is contained in:
i-robot 2021-07-01 02:01:54 +00:00 committed by Gitee
parent bb76628b70 3e4619c105
commit 81ae4d950c
2 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mindspore/lite/src/train/train_export.cc
View File

@ -19,6 +19,7 @@
#include <sys/stat.h>
#include <fstream>
#include <utility>
#include <functional>
#include <map>
#include <set>
#include "schema/inner/model_generated.h"
@ -375,7 +376,7 @@ int TrainExport::ExportNet(const std::vector<mindspore::kernel::LiteKernel *> &k
return RET_ERROR;
}
if (out_set.find(remap_[id]) == out_set.end()) {
if ((tensorT->nodeType == NodeType_ValueNode) && (tensorT->data.size() == 0)) {
if (IsInputTensor(*tensorT)) {
meta_graph_->inputIndex.push_back(remap_[id]);
if (!meta_graph_->subGraph.empty()) {
meta_graph_->subGraph[0]->inputIndices.push_back(remap_[id]);
@ -413,6 +414,11 @@ int TrainExport::ExportInit(const std::string model_name, std::string version) {
int TrainExport::SaveToFile() { return Storage::Save(*meta_graph_, file_name_); }
int TrainExport::IsInputTensor(const schema::TensorT &t) {
int total_dims = std::accumulate(t.dims.begin(), t.dims.end(), 1, std::multiplies<int>());
return ((t.nodeType == NodeType_ValueNode) && (t.data.size() == 0) && (total_dims != 0));
}
TrainExport::~TrainExport() { delete meta_graph_; }
} // namespace lite

1
mindspore/lite/src/train/train_export.h
View File

@ -63,6 +63,7 @@ class TrainExport {
std::unique_ptr<schema::CNodeT> CreateCNode(const mindspore::kernel::LiteKernel *kernel,
std::vector<uint32_t> inputIndex, std::vector<uint32_t> outputIndex,
const Model *model);
int IsInputTensor(const schema::TensorT &t);
int CreateAndAddCNode(const mindspore::kernel::LiteKernel *kernel, std::vector<uint32_t> inputIndex,
std::vector<uint32_t> outputIndex, const Model *model);
std::unique_ptr<schema::CNodeT> CreateTransformNode(std::vector<uint32_t> inputIndex,