create multi make tuple for nested tuple output

This commit is contained in:
lvliang 2021-05-13 18:11:59 +08:00 committed by chujinjin
parent 2d512ab991
commit 32da2cea42
5 changed files with 21 additions and 25 deletions

View File

@ -108,11 +108,6 @@ void SyncDeviceInfoToValueNode(const ValueNodePtr &value_node, std::vector<std::
std::vector<tensor::TensorPtr> tensors;
TensorValueToTensor(value, &tensors);
if (!tensors.empty()) {
if (tensors.size() != AnfAlgo::GetOutputTensorNum(value_node)) {
MS_LOG(EXCEPTION) << "The size of tensors converted from value [" << tensors.size()
<< "] is not equal to output size of value node [" << AnfAlgo::GetOutputTensorNum(value_node)
<< "]";
}
device_formats->clear();
device_types->clear();
for (const auto &tensor : tensors) {

View File

@ -534,9 +534,9 @@ PynativeAdjointPtr KPynativeCellImpl::ForgeMakeSequenceAdjoint(const CNodePtr &c
// () or [] is not supported yet.
if (cnode->size() <= 1) {
MS_LOG(DEBUG) << "MakeTuple/MakeList CNode is empty Tuple/List, CNode: " << cnode->DebugString();
static auto empty_tuple = MakeValue(std::vector<ValuePtr>{});
static auto dummy_adjoint =
std::make_shared<PynativeAdjoint>(tape_, ValuePtrList{}, empty_tuple, FuncGraphPtr(nullptr));
auto empty_tuple = MakeValue(std::vector<ValuePtr>{});
auto dummy_adjoint =
std::make_shared<PynativeAdjoint>(FuncGraphPtr(nullptr), ValuePtrList{}, empty_tuple, FuncGraphPtr(nullptr));
anfnode_to_adjoin_[cnode] = dummy_adjoint;
cnode->set_stop_gradient(true);
return dummy_adjoint;

View File

@ -2050,24 +2050,24 @@ void GradExecutor::SetTupleItemArgsToGraphInfoMap(const FuncGraphPtr &g, const p
}
}
void GradExecutor::SetMakeTupleAsOutputNode(const std::string &cell_id, const FuncGraphPtr &curr_g,
const py::object &out) {
void GradExecutor::CreateMakeTupleNodeForMultiOut(const std::string &cell_id, const FuncGraphPtr &curr_g,
const py::object &out) {
MS_EXCEPTION_IF_NULL(curr_g);
if (!(py::isinstance<py::tuple>(out) || py::isinstance<py::list>(out))) {
MS_LOG(EXCEPTION) << "The out of top cell should be tuple or list when set maketuple as output node";
}
auto tuple = out.cast<py::tuple>();
auto tuple_size = static_cast<int64_t>(tuple.size());
auto out_tuple = out.cast<py::tuple>();
auto out_tuple_size = static_cast<int64_t>(out_tuple.size());
// get input node and value
std::vector<AnfNodePtr> inputs{NewValueNode(prim::kPrimMakeTuple)};
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]));
for (int64_t i = 0; i < out_tuple_size; i++) {
inputs.emplace_back(GetInput(out_tuple[i], false));
input_args.emplace_back(parse::data_converter::PyDataToValue(out_tuple[i]));
}
auto cnode = curr_g_->NewCNode(inputs);
MS_LOG(DEBUG) << "Tuple output node info " << cnode->DebugString();
// record node info in graph map
auto out_id = GetId(out);
SetTupleArgsToGraphInfoMap(curr_g_, out, cnode);
@ -2079,7 +2079,6 @@ void GradExecutor::SetMakeTupleAsOutputNode(const std::string &cell_id, const Fu
// run ad for maketuple node
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();
}
void GradExecutor::EndGraphInner(py::object *ret, const py::object &cell, const py::object &out, const py::args &args) {
@ -2098,7 +2097,7 @@ void GradExecutor::EndGraphInner(py::object *ret, const py::object &cell, const
MS_EXCEPTION_IF_NULL(graph_info);
if (graph_info->node_map.find(out_id) == graph_info->node_map.end()) {
if (py::isinstance<py::tuple>(out) || py::isinstance<py::list>(out)) {
SetMakeTupleAsOutputNode(cell_id, curr_g_, out);
CreateMakeTupleNodeForMultiOut(cell_id, curr_g_, out);
} else {
MS_LOG(DEBUG) << "Set ValueNode as output for graph, out id: " << out_id;
MakeValueNode(out, out_id);
@ -2124,11 +2123,7 @@ void GradExecutor::EndGraphInner(py::object *ret, const py::object &cell, const
set_fg_fn();
DumpIR("fg.ir", curr_g_);
}
// Checkout whether need to compile graph when top cell has ran finished
if (cell_id == top_cell()->cell_id()) {
CheckNeedCompileGraph();
}
// Reset grad flag and checkout whether need to compile graph when top cell has ran finished
// Reset grad flag and update output node of top cell
if (cell_stack_.empty() && cell_id == top_cell()->cell_id()) {
MS_LOG(DEBUG) << "Cur top last cell " << cell_id;
set_grad_flag(false);
@ -2139,6 +2134,10 @@ void GradExecutor::EndGraphInner(py::object *ret, const py::object &cell, const
MS_EXCEPTION_IF_NULL(k_pynative_cell_ptr);
k_pynative_cell_ptr->UpdateOutputNodeOfTopCell(output_node);
}
// Checkout whether need to compile graph when top cell has ran finished
if (cell_id == top_cell()->cell_id()) {
CheckNeedCompileGraph();
}
}
void GradExecutor::DoGradForCustomBprop(const py::object &cell, const py::object &out, const py::args &args) {

View File

@ -257,7 +257,7 @@ class GradExecutor {
const std::vector<int64_t> &index) {
top_cell()->graph_info_map()[g]->node_map[id] = std::make_pair(node, index);
}
void SetMakeTupleAsOutputNode(const std::string &cell_id, const FuncGraphPtr &curr_g, const py::object &out);
void CreateMakeTupleNodeForMultiOut(const std::string &cell_id, const FuncGraphPtr &curr_g, const py::object &out);
void DoGradForCustomBprop(const py::object &cell, const py::object &out, const py::args &args);
private:

View File

@ -287,6 +287,8 @@ void TensorValueToTensor(const ValuePtr &value, std::vector<tensor::TensorPtr> *
auto tensor = element->cast<tensor::TensorPtr>();
MS_EXCEPTION_IF_NULL(tensor);
tensors->emplace_back(tensor);
} else if (element->isa<ValueTuple>()) {
TensorValueToTensor(element, tensors);
}
}
} else if (value->isa<tensor::Tensor>()) {