You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
85 lines
2.9 KiB
85 lines
2.9 KiB
/**
|
|
* Copyright 2020 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 "graph/passes/save_pass.h"
|
|
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
#include "framework/common/debug/ge_log.h"
|
|
#include "common/ge_inner_error_codes.h"
|
|
#include "graph/utils/graph_utils.h"
|
|
|
|
namespace ge {
|
|
namespace {
|
|
const char *const kSave = "Save";
|
|
const char *const kVar = "Variable";
|
|
const char *const kVarIsSave = "save_checkpoint";
|
|
const char *const kVarAttrVarIsSave = "_var_is_save";
|
|
} // namespace
|
|
|
|
Status SavePass::Run(ge::ComputeGraphPtr graph) {
|
|
GE_CHECK_NOTNULL(graph);
|
|
vector<NodePtr> front_nodes;
|
|
vector<uint8_t> out_index;
|
|
vector<NodePtr> del_nodes;
|
|
for (auto &node : graph->GetDirectNode()) {
|
|
if (node->GetType() == kSave) {
|
|
for (auto &in : node->GetAllInDataAnchors()) {
|
|
auto out_anchor = in->GetPeerOutAnchor();
|
|
if (out_anchor != nullptr) {
|
|
ge::NodePtr peer_node = out_anchor->GetOwnerNode();
|
|
if (peer_node->GetType() == kVar) {
|
|
front_nodes.emplace_back(peer_node);
|
|
out_index.emplace_back(out_anchor->GetIdx());
|
|
ge::OpDescPtr op_desc = peer_node->GetOpDesc();
|
|
GE_IF_BOOL_EXEC(!ge::AttrUtils::SetStr(op_desc, kVarAttrVarIsSave, kVarIsSave),
|
|
GELOGE(INTERNAL_ERROR, "get kVarAttrVarIsSave failed"); return INTERNAL_ERROR);
|
|
}
|
|
}
|
|
}
|
|
del_nodes.emplace_back(node);
|
|
}
|
|
}
|
|
// add output nodes for save
|
|
std::vector<std::pair<NodePtr, int32_t>> out_nodes_info{};
|
|
for (size_t i = 0; i < front_nodes.size(); i++) {
|
|
out_nodes_info.emplace_back(pair<NodePtr, int32_t>(front_nodes[i], out_index[i]));
|
|
}
|
|
graph->AppendGraphOutNodesInfo(out_nodes_info);
|
|
|
|
// delete save node
|
|
for (auto &node_ptr : del_nodes) {
|
|
auto ret = graph->RemoveNode(node_ptr);
|
|
if (ret != SUCCESS) {
|
|
GELOGE(ret, "GraphUtils::RemoveNodeWithoutRelink failed.");
|
|
return ret;
|
|
}
|
|
|
|
// update Target list
|
|
vector<NodePtr> graph_target = graph->GetGraphTargetNodesInfo();
|
|
auto iter = find(graph_target.begin(), graph_target.end(), node_ptr);
|
|
if (iter != graph_target.end()) {
|
|
GELOGI("Current node %s is as Target, remove it from target vector.", node_ptr->GetName().c_str());
|
|
graph_target.erase(iter);
|
|
graph->SetGraphTargetNodesInfo(graph_target);
|
|
}
|
|
}
|
|
|
|
return SUCCESS;
|
|
}
|
|
} // namespace ge
|