diff --git a/paddle/fluid/framework/block_desc.cc b/paddle/fluid/framework/block_desc.cc index 0dd37e7df067454ef828e0569ddf0cc99e7b73d6..996aefd0479f7e15f4f1bcd0b4356456ad84bd39 100644 --- a/paddle/fluid/framework/block_desc.cc +++ b/paddle/fluid/framework/block_desc.cc @@ -46,11 +46,25 @@ VarDesc *BlockDesc::FindVarRecursive(const std::string &name) const { if (name == kEmptyVarName) return nullptr; auto it = vars_.find(name); - if (it == vars_.end()) { - return Parent() == kNoneBlockIndex ? nullptr - : ParentBlock()->FindVarRecursive(name); + if (it != vars_.end()) { + return it->second.get(); } - return it->second.get(); + + BlockDesc *tmp = ParentBlock(); + + if (tmp != nullptr) { + auto ptr = tmp->FindVarRecursive(name); + if (ptr != nullptr) { + return ptr; + } + } + + tmp = ForwardBlock(); + if (tmp != nullptr) { + return tmp->FindVarRecursive(name); + } + + return nullptr; } VarDesc &BlockDesc::FindRecursiveOrCreateVar(const std::string &name_bytes) { @@ -136,10 +150,7 @@ void BlockDesc::Flush() { } BlockDesc *BlockDesc::ParentBlock() const { - if (this->desc_->parent_idx() == kNoneBlockIndex) { - return nullptr; - } - return prog_->MutableBlock(static_cast(this->desc_->parent_idx())); + return prog_->MutableBlock(static_cast(desc_->parent_idx())); } proto::BlockDesc *BlockDesc::Proto() { @@ -186,5 +197,16 @@ void BlockDesc::ClearPBVars() { } } +void BlockDesc::SetForwardBlockID(int32_t forward_block_id) { + PADDLE_ENFORCE(!desc_->has_forward_block_idx(), + "Parent block ID has been set to %d. Cannot set to %d", + desc_->forward_block_idx(), forward_block_id); + desc_->set_forward_block_idx(forward_block_id); +} + +BlockDesc *BlockDesc::ForwardBlock() const { + return prog_->MutableBlock(static_cast(desc_->forward_block_idx())); +} + } // namespace framework } // namespace paddle diff --git a/paddle/fluid/framework/block_desc.h b/paddle/fluid/framework/block_desc.h index 4e2b03e245f34eb26ad71d865a3fe99ba5f85dc9..8345934a71c7c828513cee0c718079ec6dea99fc 100644 --- a/paddle/fluid/framework/block_desc.h +++ b/paddle/fluid/framework/block_desc.h @@ -49,6 +49,8 @@ class BlockDesc { int32_t Parent() const { return desc_->parent_idx(); } + int32_t ForwardBlockID() const { return desc_->forward_block_idx(); } + VarDesc *Var(const std::string &name_bytes); VarDesc *FindVar(const std::string &name_bytes) const; @@ -73,6 +75,10 @@ class BlockDesc { BlockDesc *ParentBlock() const; + BlockDesc *ForwardBlock() const; + + void SetForwardBlockID(int32_t forward_block_id); + OpDesc *AppendOp(); void AppendAllocatedOp(std::unique_ptr &&op_desc); @@ -91,7 +97,7 @@ class BlockDesc { proto::BlockDesc *Proto(); - ProgramDesc *Program() { return this->prog_; } + ProgramDesc *Program() const { return this->prog_; } private: void ClearPBOps(); diff --git a/paddle/fluid/framework/framework.proto b/paddle/fluid/framework/framework.proto index 4eb18b4e4d685111d02387d5ab944146c9217e62..5b43f5a8a4a1c128b04ac206d387e30c55f533fe 100644 --- a/paddle/fluid/framework/framework.proto +++ b/paddle/fluid/framework/framework.proto @@ -158,6 +158,7 @@ message BlockDesc { required int32 parent_idx = 2; repeated VarDesc vars = 3; repeated OpDesc ops = 4; + optional int32 forward_block_idx = 5 [ default = -1 ]; } // Please refer to diff --git a/paddle/fluid/framework/program_desc.h b/paddle/fluid/framework/program_desc.h index 8d4b999ad2fc0924d4609415172b87ac7c6357e9..538a0372116e6f90fd2fae5f00097b8efc5dcb5c 100644 --- a/paddle/fluid/framework/program_desc.h +++ b/paddle/fluid/framework/program_desc.h @@ -38,7 +38,13 @@ class ProgramDesc { BlockDesc *AppendBlock(const BlockDesc &parent); - BlockDesc *MutableBlock(size_t idx) { return blocks_[idx].get(); } + BlockDesc *MutableBlock(size_t idx) { + if (idx == static_cast(kNoneBlockIndex)) { + return nullptr; + } else { + return blocks_[idx].get(); + } + } const BlockDesc &Block(size_t idx) const { return *blocks_[idx]; } diff --git a/paddle/fluid/operators/while_op.cc b/paddle/fluid/operators/while_op.cc index 3d5cdeda26ada94fbd8e6a7c25995aa7de93fb3d..5f51a273dd4d2e45b20e3714805297dad85196aa 100644 --- a/paddle/fluid/operators/while_op.cc +++ b/paddle/fluid/operators/while_op.cc @@ -231,7 +231,8 @@ class WhileGradOpDescMaker : public framework::SingleGradOpDescMaker { while_grad->SetInput(kStepScopes, Output(kStepScopes)); auto *grad_block = this->grad_block_[0]; - auto *fwd_block = grad_block->ParentBlock(); + auto *fwd_block = grad_block->ForwardBlock(); + auto *parent_block = grad_block->ParentBlock(); // Not all of IGs will be generated by inner gradient operators of while op. // Ignore IGs that is not generated by the inside block. @@ -265,8 +266,10 @@ class WhileGradOpDescMaker : public framework::SingleGradOpDescMaker { for (auto &input_name : op->InputArgumentNames()) { // If the input of Op has been recorded or is generated by the forward // block, do not make it as input again. + if (block_ins.find(input_name) != block_ins.end() || - fwd_block->FindVar(input_name) != nullptr) { + fwd_block->FindVar(input_name) != nullptr || + parent_block->FindVar(input_name) != nullptr) { continue; } extra_inputs.insert(input_name); diff --git a/paddle/fluid/pybind/protobuf.cc b/paddle/fluid/pybind/protobuf.cc index 131971099ef3febc3cfaff30e918fa74cfc6cfe4..01dc53de78b1026d965ce7da140f43a794102c45 100644 --- a/paddle/fluid/pybind/protobuf.cc +++ b/paddle/fluid/pybind/protobuf.cc @@ -155,6 +155,8 @@ void BindBlockDesc(py::module &m) { py::class_(m, "BlockDesc", "") .def_property_readonly("id", &BlockDesc::ID) .def_property_readonly("parent", &BlockDesc::Parent) + .def("get_forward_block_idx", &BlockDesc::ForwardBlockID) + .def("set_forward_block_idx", &BlockDesc::SetForwardBlockID) .def("append_op", &BlockDesc::AppendOp, py::return_value_policy::reference) .def("prepend_op", &BlockDesc::PrependOp, diff --git a/python/paddle/v2/fluid/backward.py b/python/paddle/v2/fluid/backward.py index 33ff43f69304ddd9330c61114dba85994b5f1bdd..ba27aaa24601bd72bcdbd064242ea2b1c345340c 100644 --- a/python/paddle/v2/fluid/backward.py +++ b/python/paddle/v2/fluid/backward.py @@ -298,7 +298,8 @@ def _append_backward_ops_(block, # If the op has its own sub-block, deal with the sub-block first if op.has_attr("sub_block"): sub_block = program.block(op.block_attr("sub_block")) - grad_sub_block = program.create_block(parent_idx=sub_block.idx) + grad_sub_block = program.create_block() + grad_sub_block.set_forward_block_idx(sub_block.idx) cb = _callback_lookup_(op) if cb is not None: if callbacks is None: @@ -310,6 +311,8 @@ def _append_backward_ops_(block, else: _append_backward_ops_(sub_block, sub_block.ops, grad_sub_block, no_grad_dict, grad_to_var, callbacks) + + program.rollback() grad_sub_block_list.append(grad_sub_block.desc) # Getting op's corresponding grad_op diff --git a/python/paddle/v2/fluid/framework.py b/python/paddle/v2/fluid/framework.py index 0e11709296a4fc7121611c1f9928314810f35783..7ec04013c91a9282e9f533c62be892d902d830a7 100644 --- a/python/paddle/v2/fluid/framework.py +++ b/python/paddle/v2/fluid/framework.py @@ -678,6 +678,13 @@ class Block(object): def parent_idx(self): return self.desc.parent + @property + def forward_block_idx(self): + return self.desc.get_forward_block_idx() + + def set_forward_block_idx(self, idx): + self.desc.set_forward_block_idx(idx) + @property def idx(self): return self.desc.id @@ -695,11 +702,22 @@ class Block(object): return self.var(name) else: if self.idx == 0: - raise ValueError("var %s is not in block(%d) nor its parents." % - name, self.idx) + raise ValueError( + "var {0} is not in block({1}) nor its parents.".format( + name, self.idx)) else: - parent_block = self.program.block(self.parent_idx) - return parent_block.var_recursive(name) + # DFS + try: + parent_block = self.program.block(self.parent_idx) + return parent_block.var_recursive(name) + except ValueError: + fwd_block = self.program.block( + self.forward_block_idx + ) if self.forward_block_idx != -1 else None + if fwd_block is not None: + return fwd_block.var_recursive(name) + else: + raise def all_parameters(self): return list(self.iter_parameters())