fix_bug_of_index_out_of_tuple

This commit is contained in:
lvliang 2021-03-25 19:30:00 +08:00 committed by chujinjin
parent 9ebd20a092
commit b54b9d22d0
4 changed files with 147 additions and 78 deletions

View File

@ -257,6 +257,8 @@ class KPynativeCellImpl : public KPynativeCell {
FuncGraphPtr BuildBpropCutFuncGraph(const PrimitivePtr &prim, const CNodePtr &cnode);
// Back propagate for MakeList or MakeTuple is generated from MetaFuncGraph.
FuncGraphPtr BuildMakeSequenceBprop(const PrimitivePtr &prim, const CNodePtr &cnode);
// Set return node according to grad flag
void SetReturnNodeByGradFlag(const AnfNodePtrList &weights, bool grad_inputs, bool grad_weights);
};
using KPynativeCellImplPtr = std::shared_ptr<KPynativeCellImpl>;
@ -294,67 +296,40 @@ FuncGraphPtr KPynativeCellImpl::Finish(const AnfNodePtrList &weights, bool grad_
auto sens_node = BuildOnesLikeValue(tape_, last_node_adjoint_iter->second->out());
last_node_adjoint_iter->second->AccumulateDout(sens_node);
}
for (size_t i = 0; i < weights.size(); ++i) {
TraceGuard trace_guard(std::make_shared<TraceCopy>(weights[i]->debug_info()));
// Add weights parameter
for (const auto &weight : weights) {
TraceGuard trace_guard(std::make_shared<TraceCopy>(weight->debug_info()));
auto p = tape_->add_parameter();
auto input_w = weights[i]->cast<ParameterPtr>();
auto input_w = weight->cast<ParameterPtr>();
MS_EXCEPTION_IF_NULL(input_w);
p->set_default_param(input_w->default_param());
}
// BackPropagate sensitivity;
BackPropagate();
// Return the gradient;
AnfNodePtrList node_list{NewValueNode(prim::kPrimMakeTuple)};
if (grad_inputs) {
for (auto input : cell_inputs_) {
auto input_adjoint_iter = anfnode_to_adjoin_.find(input);
if (input_adjoint_iter == anfnode_to_adjoin_.end()) {
// If input is not used in the network, just return zeros_like() as dout;
MS_LOG(WARNING) << "Input is not used in network, input: " << input->ToString();
auto dout = BuildZerosLikeNode(tape_, input);
node_list.push_back(dout);
} else {
node_list.push_back(input_adjoint_iter->second->RealDout());
}
}
}
if (grad_weights) {
for (auto weight : weights) {
auto input_adjoint_iter = anfnode_to_adjoin_.find(weight);
if (input_adjoint_iter == anfnode_to_adjoin_.end()) {
// If weight is not used in the network, just return zeros_like() as dout;
MS_LOG(WARNING) << "Weight is not used in network, weight: " << weight->ToString();
auto input_w = weight->cast<ParameterPtr>();
MS_EXCEPTION_IF_NULL(input_w);
auto default_param = input_w->default_param();
MS_EXCEPTION_IF_NULL(default_param);
auto dout = BuildZerosLikeValue(tape_, default_param);
node_list.push_back(dout);
} else {
node_list.push_back(input_adjoint_iter->second->RealDout());
}
}
}
auto tape_output = tape_->NewCNode(node_list);
tape_->set_output(tape_output);
SetReturnNodeByGradFlag(weights, grad_inputs, grad_weights);
// Replace AnfNode with parameter of tape_;
auto mng = MakeManager({tape_}, false);
auto tr = mng->Transact();
const auto &parameters = tape_->parameters();
for (size_t i = 0; i < cell_inputs_.size(); ++i) {
auto cell_inputs_size = cell_inputs_.size();
for (size_t i = 0; i < cell_inputs_size; ++i) {
tr.Replace(cell_inputs_[i], parameters[i]);
}
// (Inputs, sens, weights)
size_t weight_offset = cell_inputs_.size() + 1;
// (Inputs, sens, weights) or (Inputs, weights)
size_t weight_offset = cell_inputs_size;
if (has_sens_arg) {
weight_offset = weight_offset + 1;
}
for (size_t i = 0; i < weights.size(); ++i) {
tr.Replace(weights[i], parameters[weight_offset + i]);
}
tr.Commit();
DumpIR("before_final_opt.ir", tape_);
if (MsContext::GetInstance()->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG)) {
DumpIR("before_final_opt.ir", tape_);
}
return tape_;
}
@ -769,5 +744,65 @@ FuncGraphPtr KPynativeCellImpl::BuildMakeSequenceBprop(const PrimitivePtr &prim,
bprop_func_graph_cache[key] = b;
return b;
}
void KPynativeCellImpl::SetReturnNodeByGradFlag(const AnfNodePtrList &weights, bool grad_inputs, bool grad_weights) {
AnfNodePtrList grad_inputs_list{NewValueNode(prim::kPrimMakeTuple)};
if (grad_inputs) {
for (const auto &input : cell_inputs_) {
MS_EXCEPTION_IF_NULL(input);
auto input_adjoint_iter = anfnode_to_adjoin_.find(input);
if (input_adjoint_iter == anfnode_to_adjoin_.end()) {
// If input is not used in the network, just return zeros_like() as dout;
MS_LOG(WARNING) << "Input is not used in network, input: " << input->ToString();
auto dout = BuildZerosLikeNode(tape_, input);
grad_inputs_list.push_back(dout);
} else {
grad_inputs_list.push_back(input_adjoint_iter->second->RealDout());
}
}
}
AnfNodePtrList grad_weights_list{NewValueNode(prim::kPrimMakeTuple)};
if (grad_weights) {
for (const auto &weight : weights) {
MS_EXCEPTION_IF_NULL(weight);
auto input_adjoint_iter = anfnode_to_adjoin_.find(weight);
if (input_adjoint_iter == anfnode_to_adjoin_.end()) {
// If weight is not used in the network, just return zeros_like() as dout;
MS_LOG(WARNING) << "Weight is not used in network, weight: " << weight->ToString();
auto input_w = weight->cast<ParameterPtr>();
MS_EXCEPTION_IF_NULL(input_w);
auto default_param = input_w->default_param();
MS_EXCEPTION_IF_NULL(default_param);
auto dout = BuildZerosLikeValue(tape_, default_param);
grad_weights_list.push_back(dout);
} else {
grad_weights_list.push_back(input_adjoint_iter->second->RealDout());
}
}
}
AnfNodePtr tape_output;
if (grad_inputs && grad_weights) {
tape_output = tape_->NewCNode(
{NewValueNode(prim::kPrimMakeTuple), tape_->NewCNode(grad_inputs_list), tape_->NewCNode(grad_weights_list)});
} else if (grad_inputs) {
tape_output = tape_->NewCNode(grad_inputs_list);
} else if (grad_weights) {
tape_output = tape_->NewCNode(grad_weights_list);
} else if (cell_inputs_.empty()) {
tape_output = tape_->NewCNode(grad_inputs_list);
} else {
auto input_adjoint_iter = anfnode_to_adjoin_.find(cell_inputs_[0]);
if (input_adjoint_iter == anfnode_to_adjoin_.end()) {
// If input is not used in the network, just return zeros_like() as dout;
MS_LOG(WARNING) << "Input is not used in network, input: " << cell_inputs_[0]->ToString();
tape_output = BuildZerosLikeNode(tape_, cell_inputs_[0]);
} else {
tape_output = input_adjoint_iter->second->RealDout();
}
}
tape_->set_output(tape_output);
}
} // namespace ad
} // namespace mindspore

View File

@ -75,6 +75,24 @@ bool SimplifyDataStructuresPass(const ResourcePtr &res) {
return true;
}
bool TransformTopGraphPass(const ResourcePtr &res) {
if (res->func_graph() == nullptr) {
MS_LOG(EXCEPTION) << "Transform top graph error.";
}
FuncGraphPtr func_graph = res->func_graph();
if (opt::FuncGraphHasTupleInput(func_graph)) {
opt::GraphTupleParamTransform graph_trans;
func_graph = graph_trans(func_graph, res->manager());
res->set_func_graph(func_graph);
AbstractBasePtrList abs_spec_list;
auto &params = func_graph->parameters();
std::transform(params.begin(), params.end(), std::back_inserter(abs_spec_list),
[](AnfNodePtr node) { return node->abstract(); });
res->set_args_spec(abs_spec_list);
}
return true;
}
bool CleanAfterOptAPass(const ResourcePtr &res) {
MS_EXCEPTION_IF_NULL(res->func_graph());
@ -162,6 +180,9 @@ FuncGraphPtr PrimBpOptPassStep2(const opt::irpass::OptimizeIRPassLib &irpass, co
FuncGraphPtr BpropGraphFinalOptPass(const ResourcePtr &res) {
MS_EXCEPTION_IF_NULL(res);
if (!TransformTopGraphPass(res)) {
MS_LOG(EXCEPTION) << "Run TransformTopGraphPass failed";
}
opt::irpass::OptimizeIRPassLib irpass;
opt::OptPassConfig bg_final_opt_ = opt::OptPassConfig({
@ -170,7 +191,6 @@ FuncGraphPtr BpropGraphFinalOptPass(const ResourcePtr &res) {
irpass.depend_value_elim_,
irpass.reshape_eliminate_,
});
OptPassGroupMap map({{"ad_final_opt_", bg_final_opt_}});
auto bprop_graph_final_opt = opt::Optimizer::MakeOptimizer("bprop_graph_final_opt", res, map);
@ -178,6 +198,7 @@ FuncGraphPtr BpropGraphFinalOptPass(const ResourcePtr &res) {
WITH(MsProfile::GetProfile()->Step("bprop_graph_final_opt"))[&bprop_graph_final_opt, &func_graph]() {
func_graph = bprop_graph_final_opt->step(func_graph, true);
};
return func_graph;
}
@ -560,24 +581,6 @@ bool CconvPass(const ResourcePtr &res) {
return true;
}
bool TransformTopGraphPass(const ResourcePtr &res) {
if (res->func_graph() == nullptr) {
MS_LOG(EXCEPTION) << "Transform top graph error.";
}
FuncGraphPtr func_graph = res->func_graph();
if (opt::FuncGraphHasTupleInput(func_graph)) {
opt::GraphTupleParamTransform graph_trans;
func_graph = graph_trans(func_graph, res->manager());
res->set_func_graph(func_graph);
AbstractBasePtrList abs_spec_list;
auto &params = func_graph->parameters();
std::transform(params.begin(), params.end(), std::back_inserter(abs_spec_list),
[](AnfNodePtr node) { return node->abstract(); });
res->set_args_spec(abs_spec_list);
}
return true;
}
bool PipelineSplitPass(const ResourcePtr &res) { return PipelineSplit(res); }
void UpdateFuncGraphParameter(const FuncGraphPtr &func_graph) {

View File

@ -679,18 +679,21 @@ bool TopCellInfo::IsSubCell(const std::string &cell_id) const {
}
void TopCellInfo::clear() {
MS_LOG(DEBUG) << "Clear top cell info. Cell id " << cell_id_;
op_num_ = 0;
is_dynamic_ = false;
vm_compiled_ = false;
is_init_kpynative_ = false;
forward_already_run_ = false;
op_num_ = 0;
input_args_id_.clear();
all_op_info_.clear();
k_pynative_cell_ptr_ = nullptr;
if (resource_ != nullptr) {
resource_->Clean();
resource_ = nullptr;
}
df_builder_ = nullptr;
resource_->Clean();
resource_ = nullptr;
k_pynative_cell_ptr_ = nullptr;
graph_info_map_.clear();
sub_cell_list_.clear();
op_info_with_tensor_id_.clear();
@ -1308,7 +1311,8 @@ void GradExecutor::UpdateForwardTensorInfoInBpropGraph(const OpExecInfoPtr &op_e
}
MS_EXCEPTION_IF_NULL(top_cell_);
MS_EXCEPTION_IF_NULL(op_exec_info);
MS_LOG(DEBUG) << "Current op info: " << op_exec_info->op_info;
auto op_info = op_exec_info->op_info;
MS_LOG(DEBUG) << "Current op info: " << op_info;
// Get input tensors;
std::vector<tensor::TensorPtr> all_op_tensors;
for (size_t i = 0; i < op_exec_info->op_inputs.size(); ++i) {
@ -1318,7 +1322,7 @@ void GradExecutor::UpdateForwardTensorInfoInBpropGraph(const OpExecInfoPtr &op_e
TensorValueToTensor(parse::data_converter::PyDataToValue(out_real), &all_op_tensors);
// Save all tensors info of current op
if (need_construct_graph()) {
SaveOpInfo(top_cell_, op_exec_info->op_info, all_op_tensors);
SaveOpInfo(top_cell_, op_info, all_op_tensors);
}
if (already_run_top_cell_.find(top_cell_->cell_id()) == already_run_top_cell_.end()) {
@ -1332,7 +1336,11 @@ void GradExecutor::UpdateForwardTensorInfoInBpropGraph(const OpExecInfoPtr &op_e
// Non-first run
const auto &pre_top_cell = already_run_top_cell_.at(top_cell_->cell_id());
MS_EXCEPTION_IF_NULL(pre_top_cell);
const auto &pre_op_tensor_id = pre_top_cell->op_info_with_tensor_id().at(op_exec_info->op_info);
if (pre_top_cell->op_info_with_tensor_id().find(op_info) == pre_top_cell->op_info_with_tensor_id().end()) {
MS_LOG(EXCEPTION) << "Can not find op info " << op_info << " in op info with tensor id map. Top cell "
<< top_cell_->cell_id();
}
const auto &pre_op_tensor_id = pre_top_cell->op_info_with_tensor_id().at(op_info);
if (pre_op_tensor_id.size() != all_op_tensors.size()) {
MS_LOG(EXCEPTION) << "The size of pre op tensor id: " << pre_op_tensor_id.size()
<< " is not equal to the size of all tensors of current op " << all_op_tensors.size();
@ -1692,13 +1700,20 @@ void GradExecutor::ClearCellRes(const std::string &cell_id) {
it->clear();
}
top_cell_list_.clear();
already_run_top_cell_.clear();
MS_LOG(DEBUG) << "Clear all cell resources";
return;
}
for (auto it = top_cell_list_.begin(); it != top_cell_list_.end();) {
if (IsCellObjIdEq(cell_id, (*it)->cell_id())) {
auto top_cell_id = (*it)->cell_id();
if (IsCellObjIdEq(cell_id, top_cell_id)) {
(*it)->clear();
it = top_cell_list_.erase(it);
break;
if (already_run_top_cell_.find(top_cell_id) != already_run_top_cell_.end()) {
(void)already_run_top_cell_.erase(top_cell_id);
}
MS_LOG(ERROR) << "Clear top cell resource. Top cell id " << top_cell_id;
continue;
}
it++;
}
@ -1898,14 +1913,19 @@ void GradExecutor::EndGraphInner(py::object *ret, const py::object &cell, const
auto tuple = out.cast<py::tuple>();
auto tuple_size = static_cast<int64_t>(tuple.size());
ValuePtrList input_args;
std::vector<AnfNodePtr> inputs;
inputs.emplace_back(NewValueNode(prim::kPrimMakeTuple));
for (int64_t i = 0; i < tuple_size; i++) {
inputs.emplace_back(GetInput(tuple[i], false));
input_args.emplace_back(parse::data_converter::PyDataToValue(tuple[i]));
}
auto cnode = curr_g_->NewCNode(inputs);
SetTupleArgsToGraphInfoMap(curr_g_, out, cnode);
SetNodeMapInGraphInfoMap(curr_g_, out_id, cnode);
ValuePtr out_value = parse::data_converter::PyDataToValue(out);
ad::GradPynativeOp(top_cell()->k_pynative_cell_ptr(), cnode, input_args, out_value);
MS_LOG(DEBUG) << "Tuple output node info " << cnode->DebugString();
} else {
MS_LOG(DEBUG) << "Set ValueNode as output for graph, out id: " << out_id;
MakeValueNode(out, out_id);
@ -2017,7 +2037,7 @@ void GradExecutor::GradNetInner(py::object *ret, const GradOperationPtr &grad, c
MS_LOG(DEBUG) << "df_builder ptr " << df_builder.get() << " resource ptr " << resource.get();
// Get params(weights) require derivative
auto w_args = GetWeightsArgs(weights, df_builder);
auto w_args = GetWeightsArgs(grad, weights, df_builder);
// Get bprop graph of top cell
auto bprop_graph = GetBpropGraph(grad, w_args, size, args);
resource->set_func_graph(bprop_graph);
@ -2037,11 +2057,18 @@ void GradExecutor::GradNetInner(py::object *ret, const GradOperationPtr &grad, c
resource->Clean();
}
std::vector<AnfNodePtr> GradExecutor::GetWeightsArgs(const py::object &weights, const FuncGraphPtr &df_builder) {
if (!py::hasattr(weights, "__parameter_tuple__")) {
std::vector<AnfNodePtr> GradExecutor::GetWeightsArgs(const GradOperationPtr &grad, const py::object &weights,
const FuncGraphPtr &df_builder) {
MS_EXCEPTION_IF_NULL(grad);
MS_EXCEPTION_IF_NULL(df_builder);
if (!grad->get_by_list_ && py::isinstance<py::none>(weights)) {
MS_LOG(DEBUG) << "The input weight is None when run GradNetInner. Return parameters of df_builder directly";
return df_builder->parameters();
} else if (!py::hasattr(weights, "__parameter_tuple__")) {
MS_LOG(DEBUG) << "No paramter_tuple get";
return {};
}
auto tuple = weights.cast<py::tuple>();
MS_LOG(DEBUG) << "Get weights tuple size " << tuple.size();
std::vector<AnfNodePtr> w_args;
@ -2111,8 +2138,8 @@ FuncGraphPtr GradExecutor::GetBpropGraph(const GradOperationPtr &grad, const std
MS_EXCEPTION_IF_NULL(grad);
auto k_pynative_cell_ptr = top_cell()->k_pynative_cell_ptr();
MS_EXCEPTION_IF_NULL(k_pynative_cell_ptr);
auto bprop_graph = ad::GradPynativeCellEnd(k_pynative_cell_ptr, weights, grad->get_all_, grad->get_by_list_,
grad->sens_param_);
auto bprop_graph =
ad::GradPynativeCellEnd(k_pynative_cell_ptr, weights, grad->get_all_, grad->get_by_list_, grad->sens_param_);
MS_EXCEPTION_IF_NULL(bprop_graph);
MS_LOG(DEBUG) << "Top graph input params size " << arg_size;
@ -2130,7 +2157,10 @@ FuncGraphPtr GradExecutor::GetBpropGraph(const GradOperationPtr &grad, const std
MS_EXCEPTION_IF_NULL(manager);
manager->AddFuncGraph(bprop_graph);
auto optimized_bg = pipeline::PrimBpropOptimizer::GetPrimBpropOptimizerInst().BpropGraphFinalOpt(resource);
DumpIR("after_final_opt.ir", optimized_bg);
if (MsContext::GetInstance()->get_param<bool>(MS_CTX_SAVE_GRAPHS_FLAG)) {
DumpIR("after_final_opt.ir", optimized_bg);
}
optimized_bg->ClearAllManagerInfo();
return optimized_bg;
}

View File

@ -175,7 +175,7 @@ class GradExecutor {
bool grad_flag() const { return grad_flag_; }
void set_grad_flag(bool flag) { grad_flag_ = flag; }
bool in_grad_process() const { return in_grad_process_; }
bool in_cell_with_custom_bprop_() const {return custom_bprop_cell_count_ > 0;}
bool in_cell_with_custom_bprop_() const { return custom_bprop_cell_count_ > 0; }
AnfNodePtr GetInput(const py::object &obj, bool op_mask);
std::string GetCellId(const py::object &obj, const py::args &args);
std::stack<std::string> &cell_stack() { return cell_stack_; }
@ -230,7 +230,8 @@ class GradExecutor {
const py::args &args);
FuncGraphPtr GetBpropGraph(const GradOperationPtr &grad, const std::vector<AnfNodePtr> &weights, size_t arg_size,
const py::args &args);
std::vector<AnfNodePtr> GetWeightsArgs(const py::object &weights, const FuncGraphPtr &df_builder);
std::vector<AnfNodePtr> GetWeightsArgs(const GradOperationPtr &grad, const py::object &weights,
const FuncGraphPtr &df_builder);
abstract::AbstractBasePtrList GetArgsSpec(const py::args &args, const FuncGraphPtr &bprop_graph);
void SetTupleItemArgsToGraphInfoMap(const FuncGraphPtr &g, const py::object &id, const AnfNodePtr &node,
const std::vector<int64_t> &index_sequence, bool is_param = false);