Add backward and optimizer operator dependency pass. (#17746)
gongweibao
6 years ago
committed by
GitHub
parent
4cb7d32c9b
commit
fbbdc9ccad
@ -0,0 +1,223 @@
|
||||
// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// 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 <algorithm>
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "paddle/fluid/framework/details/container_cast.h"
|
||||
#include "paddle/fluid/framework/details/multi_devices_helper.h"
|
||||
#include "paddle/fluid/framework/details/op_handle_base.h"
|
||||
#include "paddle/fluid/framework/ir/graph.h"
|
||||
#include "paddle/fluid/framework/ir/graph_helper.h"
|
||||
#include "paddle/fluid/framework/ir/pass.h"
|
||||
#include "paddle/fluid/framework/lod_tensor.h"
|
||||
#include "paddle/fluid/framework/op_proto_maker.h"
|
||||
#include "paddle/fluid/framework/scope.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace framework {
|
||||
namespace ir {
|
||||
|
||||
class BackWardOpDepsPass : public ir::Pass {
|
||||
protected:
|
||||
void AddDep(ir::Graph* graph, details::OpHandleBase* l,
|
||||
details::OpHandleBase* r) const {
|
||||
auto* dep_var = new details::DummyVarHandle(graph->CreateControlDepVar());
|
||||
graph->Get<details::GraphDepVars>(details::kGraphDepVars).emplace(dep_var);
|
||||
l->AddOutput(dep_var);
|
||||
r->AddInput(dep_var);
|
||||
VLOG(10) << "add deps:" << l->DebugString() << " and " << r->DebugString();
|
||||
}
|
||||
|
||||
void ApplyImpl(ir::Graph* graph) const override {
|
||||
// NOTE: The operator nodes should be in topology order.
|
||||
std::vector<details::OpHandleBase*> backward_op_handles;
|
||||
std::vector<details::OpHandleBase*> all_opt_handles;
|
||||
details::ParamsAndGrads params_grads;
|
||||
std::vector<ir::Node*> topo_nodes = ir::TopologySortOperations(*graph);
|
||||
for (auto& node : topo_nodes) {
|
||||
if (!node->Op()) continue;
|
||||
|
||||
GetBackWardOpHandles(node, &backward_op_handles, ¶ms_grads);
|
||||
GetOptimizerOpHandles(node, &all_opt_handles);
|
||||
}
|
||||
|
||||
VLOG(10) << "backward_op_handles size:" << backward_op_handles.size()
|
||||
<< ", opt_handles size:" << all_opt_handles.size();
|
||||
|
||||
if (backward_op_handles.size() <= 1 || all_opt_handles.size() <= 1) {
|
||||
VLOG(10) << "need not backward_op_deps_pass";
|
||||
return;
|
||||
}
|
||||
|
||||
std::vector<details::OpHandleBase*> opt_handles;
|
||||
GetOptimizerHandlesRoot(all_opt_handles, &opt_handles, params_grads);
|
||||
|
||||
if (opt_handles.size() <= 1) {
|
||||
VLOG(10) << "need not backward_op_deps_pass";
|
||||
return;
|
||||
}
|
||||
|
||||
VLOG(10) << "add optimize deps";
|
||||
for (size_t i = 1; i < opt_handles.size(); ++i) {
|
||||
AddDep(graph, opt_handles[i - 1], opt_handles[i]);
|
||||
}
|
||||
|
||||
VLOG(10) << "add deps between backward and optimze:";
|
||||
AddDep(graph, backward_op_handles[backward_op_handles.size() - 1],
|
||||
opt_handles[0]);
|
||||
}
|
||||
|
||||
/*
|
||||
* When the backward ophandles complete, the optimizer ophandle's inputs var
|
||||
* are ready.Since the optimizer ophandles can be seen as graphs which each of
|
||||
* them doesn't connect to each other, they can run parallelly or by a
|
||||
* specified order, such as by the grads generated order. This function will
|
||||
* get these graphs' root.
|
||||
*/
|
||||
void GetOptimizerHandlesRoot(
|
||||
const std::vector<details::OpHandleBase*>& ops,
|
||||
std::vector<details::OpHandleBase*>* result,
|
||||
const details::ParamsAndGrads& params_grads) const {
|
||||
std::unordered_set<details::OpHandleBase*> visit;
|
||||
for (auto op : ops) {
|
||||
if (visit.find(op) != visit.end()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
VLOG(10) << "visiting all_opt_handles:" << op->DebugString();
|
||||
|
||||
result->emplace_back(op);
|
||||
visit.insert(op);
|
||||
VisitChildrens(op, &visit);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < result->size(); i++) {
|
||||
VLOG(10) << "get potential head op:" << (*result)[i]->DebugString();
|
||||
}
|
||||
|
||||
// sort by param_grad order
|
||||
std::unordered_map<std::string, int> pg_order;
|
||||
int order = 0;
|
||||
for (auto& p_g : params_grads) {
|
||||
pg_order[p_g.second] = order++;
|
||||
}
|
||||
|
||||
std::vector<std::pair<details::OpHandleBase*, int>> op_handles;
|
||||
for (auto op : *result) {
|
||||
int order = 0;
|
||||
for (auto input : op->Inputs()) {
|
||||
if (dynamic_cast<details::VarHandle*>(input) == nullptr) continue;
|
||||
|
||||
if (pg_order.find(input->Name()) == pg_order.end()) {
|
||||
VLOG(10) << "not find input " << input->Name() << " in grad";
|
||||
continue;
|
||||
}
|
||||
|
||||
if (order < pg_order.at(input->Name())) {
|
||||
order = pg_order.at(input->Name());
|
||||
}
|
||||
}
|
||||
op_handles.emplace_back(std::make_pair(op, order));
|
||||
}
|
||||
|
||||
sort(op_handles.begin(), op_handles.end(),
|
||||
[](const std::pair<details::OpHandleBase*, int>& left,
|
||||
const std::pair<details::OpHandleBase*, int>& right) -> bool {
|
||||
return left.second < right.second;
|
||||
});
|
||||
|
||||
result->clear();
|
||||
for (auto p : op_handles) {
|
||||
result->emplace_back(p.first);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < result->size(); i++) {
|
||||
VLOG(10) << "get head op:" << (*result)[i]->DebugString();
|
||||
}
|
||||
}
|
||||
|
||||
void VisitChildrens(details::OpHandleBase* op,
|
||||
std::unordered_set<details::OpHandleBase*>* visit) const {
|
||||
for (auto out : op->Outputs()) {
|
||||
for (auto* pending_op : out->PendingOps()) {
|
||||
if (visit->find(pending_op) != visit->end()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
VLOG(10) << "visiting:" << pending_op->DebugString();
|
||||
|
||||
visit->insert(pending_op);
|
||||
VisitChildrens(pending_op, visit);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void GetBackWardOpHandles(
|
||||
ir::Node* node, std::vector<details::OpHandleBase*>* backward_op_handles,
|
||||
details::ParamsAndGrads* params_grads) const {
|
||||
try {
|
||||
bool is_bk_op =
|
||||
static_cast<bool>(boost::get<int>(node->Op()->GetAttr(
|
||||
OpProtoAndCheckerMaker::OpRoleAttrName())) &
|
||||
static_cast<int>(OpRole::kBackward));
|
||||
if (!is_bk_op) return;
|
||||
|
||||
// Currently, we assume that once gradient is generated, it can be
|
||||
// broadcast, and each gradient is only broadcast once.
|
||||
auto backward_vars =
|
||||
boost::get<std::vector<std::string>>(node->Op()->GetNullableAttr(
|
||||
OpProtoAndCheckerMaker::OpRoleVarAttrName()));
|
||||
PADDLE_ENFORCE_EQ(backward_vars.size() % 2, static_cast<size_t>(0));
|
||||
PADDLE_ENFORCE(node->IsWrappedBy<details::OpHandleBase>());
|
||||
|
||||
backward_op_handles->emplace_back(
|
||||
&node->Wrapper<details::OpHandleBase>());
|
||||
|
||||
for (size_t i = 0; i < backward_vars.size(); i += 2) {
|
||||
VLOG(10) << "Trainable parameter: " << backward_vars[i]
|
||||
<< ", gradient: " << backward_vars[i + 1];
|
||||
|
||||
params_grads->emplace_back(std::make_pair(
|
||||
backward_vars[i] /*param*/, backward_vars[i + 1] /*grad*/));
|
||||
}
|
||||
} catch (boost::bad_get e) {
|
||||
}
|
||||
}
|
||||
|
||||
void GetOptimizerOpHandles(
|
||||
ir::Node* node, std::vector<details::OpHandleBase*>* opt_handles) const {
|
||||
try {
|
||||
bool is_opt_op =
|
||||
static_cast<bool>(boost::get<int>(node->Op()->GetAttr(
|
||||
OpProtoAndCheckerMaker::OpRoleAttrName())) &
|
||||
static_cast<int>(OpRole::kOptimize));
|
||||
if (!is_opt_op) return;
|
||||
|
||||
opt_handles->emplace_back(&node->Wrapper<details::OpHandleBase>());
|
||||
} catch (boost::bad_get e) {
|
||||
}
|
||||
}
|
||||
};
|
||||
} // namespace ir
|
||||
} // namespace framework
|
||||
} // namespace paddle
|
||||
|
||||
REGISTER_PASS(backward_optimizer_op_deps_pass,
|
||||
paddle::framework::ir::BackWardOpDepsPass);
|
||||
Loading…
Reference in new issue