forked from huawei/mindspore2022
102 lines
3.6 KiB
C++
102 lines
3.6 KiB
C++
/**
|
|
* Copyright 2021 Huawei Technologies Co., Ltd
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#include "coder/train.h"
|
|
#include <memory>
|
|
#include <set>
|
|
#include <array>
|
|
#include <queue>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <algorithm>
|
|
#include "schema/ops_generated.h"
|
|
#include "src/common/prim_util.h"
|
|
|
|
namespace mindspore::lite::micro {
|
|
std::set<OperatorCoder *> FindInferenceOpcoders(OperatorCoder *edge) {
|
|
std::set<OperatorCoder *> subgraph;
|
|
std::queue<OperatorCoder *> to_visit;
|
|
to_visit.push(edge);
|
|
while (!to_visit.empty()) {
|
|
size_t size = to_visit.size();
|
|
for (size_t i = 0; i < size; ++i) {
|
|
OperatorCoder *curr = to_visit.front();
|
|
to_visit.pop();
|
|
if (subgraph.find(curr) != subgraph.end()) {
|
|
continue;
|
|
}
|
|
subgraph.insert(curr);
|
|
for (const auto &op : curr->input_ops()) {
|
|
to_visit.push(op);
|
|
}
|
|
}
|
|
}
|
|
auto item = subgraph.find(edge);
|
|
if (item == subgraph.end()) {
|
|
MS_LOG(ERROR) << "failed to find the edge in the subgraph";
|
|
return subgraph;
|
|
}
|
|
// erase edge operator coder from subgraph
|
|
subgraph.erase(item);
|
|
return subgraph;
|
|
}
|
|
|
|
int Train::TransformGraphForTrain(CoderContext *context, const std::vector<std::unique_ptr<OperatorCoder>> &op_coders,
|
|
int schema_version) {
|
|
if (context == nullptr) {
|
|
MS_LOG(INFO) << "input context invalid";
|
|
return RET_ERROR;
|
|
}
|
|
const std::array<int, 6> loss_types = {schema::PrimitiveType_SparseSoftmaxCrossEntropyWithLogits,
|
|
schema::PrimitiveType_BinaryCrossEntropy,
|
|
schema::PrimitiveType_SmoothL1Loss,
|
|
schema::PrimitiveType_SmoothL1LossGrad,
|
|
schema::PrimitiveType_SigmoidCrossEntropyWithLogits,
|
|
schema::PrimitiveType_SigmoidCrossEntropyWithLogitsGrad};
|
|
OperatorCoder *loss_op = nullptr;
|
|
for (const auto &opcoder : op_coders) {
|
|
const Model::Node *node = opcoder->node();
|
|
int primitive_type = GetPrimitiveType(node->primitive_, schema_version);
|
|
auto item = std::find(loss_types.begin(), loss_types.end(), primitive_type);
|
|
if (item != loss_types.end()) {
|
|
loss_op = opcoder.get();
|
|
break;
|
|
}
|
|
}
|
|
MS_CHECK_PTR(loss_op);
|
|
size_t op_num = op_coders.size();
|
|
std::vector<std::string> code_blocks = context->code_blocks();
|
|
if (op_num != code_blocks.size()) {
|
|
MS_LOG(INFO) << "the number of code blocks and op coders is not equal";
|
|
return RET_ERROR;
|
|
}
|
|
std::set<OperatorCoder *> inference_ops = FindInferenceOpcoders(loss_op);
|
|
std::vector<std::string> inferences_blocks;
|
|
std::vector<std::string> train_blocks;
|
|
for (size_t i = 0; i < op_num; ++i) {
|
|
auto &opcoder = op_coders.at(i);
|
|
std::string block = code_blocks.at(i);
|
|
if (inference_ops.find(opcoder.get()) != inference_ops.end()) {
|
|
inferences_blocks.push_back(block);
|
|
}
|
|
train_blocks.push_back(block);
|
|
}
|
|
context->set_inference_blocks(inferences_blocks);
|
|
context->set_train_blocks(train_blocks);
|
|
return RET_OK;
|
|
}
|
|
} // namespace mindspore::lite::micro
|