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.
116 lines
4.2 KiB
116 lines
4.2 KiB
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
|
|
|
|
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 "paddle/framework/grad_op_builder.h"
|
|
#include "paddle/framework/op_registry.h"
|
|
|
|
namespace paddle {
|
|
namespace framework {
|
|
|
|
OperatorBase* GradOpBuilder::Build() {
|
|
BuildOpInOutArgList();
|
|
std::string grad_op_type = OpRegistry::grad_ops().at(op_.type_);
|
|
OperatorBase* grad_op = OpRegistry::op_creators().at(grad_op_type)();
|
|
grad_op->type_ = grad_op_type;
|
|
CompleteGradOp(grad_op);
|
|
return grad_op;
|
|
}
|
|
|
|
OpInOutArg* GradOpBuilder::BuildArg(const VarProto& var,
|
|
const VarIndexMap& var_map,
|
|
const std::vector<int>& format,
|
|
InOutType type) {
|
|
int idx = var_map.at(var.name());
|
|
int begin_idx = format.empty() ? idx : format.at(idx);
|
|
int end_idx = format.empty() ? idx + 1 : format.at(idx + 1);
|
|
return new OpInOutArg(var.name(), type, !var.ignore_gradient(), begin_idx,
|
|
end_idx);
|
|
}
|
|
|
|
void GradOpBuilder::BuildOpInOutArgList() {
|
|
const OpProto& op_proto = OpRegistry::protos().at(op_.type_);
|
|
const auto& var_map = *(OpRegistry::VarIndexMaps().at(op_.type_));
|
|
const std::vector<int>& in_format =
|
|
op_.attrs_.count("input_format")
|
|
? op_.GetAttr<std::vector<int>>("input_format")
|
|
: std::vector<int>();
|
|
const std::vector<int>& out_format =
|
|
op_.attrs_.count("output_format")
|
|
? op_.GetAttr<std::vector<int>>("output_format")
|
|
: std::vector<int>();
|
|
for (const auto& var : op_proto.inputs()) {
|
|
arg_list_.emplace_back(
|
|
std::shared_ptr<OpInOutArg>(BuildArg(var, var_map, in_format, IN)));
|
|
}
|
|
for (const auto& var : op_proto.outputs()) {
|
|
arg_list_.emplace_back(
|
|
std::shared_ptr<OpInOutArg>(BuildArg(var, var_map, out_format, OUT)));
|
|
}
|
|
}
|
|
|
|
void GradOpBuilder::AddArgIntoGradOp(const OpInOutArg* arg,
|
|
std::vector<std::string>& in_out,
|
|
std::vector<int>& format,
|
|
VarIndexMap* varmap, int& idx,
|
|
bool is_grad) const {
|
|
std::string var_name = arg->proto_name_;
|
|
if (is_grad) {
|
|
var_name += OperatorBase::GRAD_VAR_SUFFIX();
|
|
}
|
|
(*varmap)[var_name] = idx++;
|
|
size_t pre_sz = in_out.size();
|
|
auto base_it = arg->type_ == IN ? op_.inputs_.begin() : op_.outputs_.begin();
|
|
std::copy(base_it + arg->begin_idx_, base_it + arg->end_idx_,
|
|
std::back_inserter(in_out));
|
|
if (is_grad) {
|
|
for (size_t i = pre_sz; i < in_out.size(); ++i) {
|
|
in_out[i] += OperatorBase::GRAD_VAR_SUFFIX();
|
|
}
|
|
}
|
|
format.push_back(in_out.size());
|
|
}
|
|
|
|
void GradOpBuilder::CompleteGradOp(OperatorBase* grad_op) const {
|
|
grad_op->attrs_ = op_.attrs_;
|
|
grad_op->attrs_.erase("input_format");
|
|
grad_op->attrs_.erase("output_format");
|
|
VarIndexMap* grad_varmap = new VarIndexMap();
|
|
int in_idx = 0;
|
|
int out_idx = 0;
|
|
std::vector<int> in_format({0});
|
|
std::vector<int> out_format({0});
|
|
for (const auto& arg : arg_list_) {
|
|
// op_'s inputs_ and outputs_
|
|
if (arg->needed_in_grad_) {
|
|
AddArgIntoGradOp(arg.get(), grad_op->inputs_, in_format, grad_varmap,
|
|
in_idx, false);
|
|
}
|
|
if (arg->type_ == IN) {
|
|
// gradients of op_'s inputs_
|
|
AddArgIntoGradOp(arg.get(), grad_op->outputs_, out_format, grad_varmap,
|
|
out_idx, true);
|
|
} else {
|
|
// gradients of op_'s outputs_
|
|
AddArgIntoGradOp(arg.get(), grad_op->inputs_, in_format, grad_varmap,
|
|
in_idx, true);
|
|
}
|
|
}
|
|
grad_op->attrs_["input_format"] = in_format;
|
|
grad_op->attrs_["output_format"] = out_format;
|
|
grad_op->in_out_idxs_.reset(grad_varmap);
|
|
}
|
|
|
|
} // namespace framework
|
|
} // namespace paddle
|