layer.cc 6.2 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
// Copyright (c) 2018 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 "paddle/fluid/imperative/layer.h"
#include <deque>
#include <limits>
#include <map>
#include <random>
#include <utility>

#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_registry.h"
24
#include "paddle/fluid/framework/operator.h"
25 26 27 28 29
#include "paddle/fluid/string/printf.h"

namespace paddle {
namespace imperative {

X
Xin Pan 已提交
30 31
std::map<int, py::object> py_funcs_;

32 33 34 35 36
using framework::Variable;

void AddTo(Variable* src, Variable* dst) {
  framework::LoDTensor* dst_tensor = dst->GetMutable<framework::LoDTensor>();
  framework::LoDTensor* src_tensor = src->GetMutable<framework::LoDTensor>();
M
minqiyang 已提交
37 38 39 40 41
  // FIXME(minqiyang): loss_grad op will pass a zero grad of label
  // ugly fix for it
  if (src_tensor->numel() == 0) {
    return;
  }
42 43 44
  PADDLE_ENFORCE(dst_tensor->numel() == src_tensor->numel(),
                 "dst_numel %lld vs. src_numel %lld", dst_tensor->numel(),
                 src_tensor->numel());
45 46 47 48 49 50 51 52 53
  float* dst_data = dst_tensor->mutable_data<float>(platform::CPUPlace());
  const float* src_data = src_tensor->data<float>();
  for (size_t i = 0; i < src_tensor->numel(); ++i) {
    dst_data[i] += src_data[i];
  }
}

class Autograd {
 public:
X
Xin Pan 已提交
54
  Autograd() {}
55 56

  void RunBackward(VarBase* var) {
57 58 59
    if (var->stop_gradient_) {
      return;
    }
60 61 62 63 64 65 66 67 68

    std::deque<OpBase*> ready;
    ready.push_back(var->pre_op_);

    std::map<OpBase*, int> dep_counts = ComputeDepCounts(var->pre_op_);

    while (!ready.empty()) {
      OpBase* ready_op = ready.front();
      ready.pop_front();
X
Xin Pan 已提交
69 70 71 72 73 74 75
      std::map<std::string, std::vector<VarBase*>> input_grads =
          ready_op->ApplyGrad();

      for (auto it : input_grads) {
        const std::vector<VarBase*>& ingrads = it.second;
        for (size_t i = 0; i < ingrads.size(); ++i) {
          if (!ingrads[i]) continue;
76 77 78
          if (ready_op->input_vars_[it.first][i]->stop_gradient_) {
            continue;
          }
X
Xin Pan 已提交
79
          OpBase* pre_op = ready_op->pre_ops_[it.first][i];
X
Xin Pan 已提交
80 81 82 83 84 85 86 87
          if (!pre_op) continue;

          dep_counts[pre_op] -= 1;
          PADDLE_ENFORCE(dep_counts[pre_op] >= 0);
          bool pre_op_ready = dep_counts[pre_op] == 0;
          if (pre_op_ready) {
            ready.push_back(pre_op);
          }
88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103
        }
      }
    }
  }

 private:
  std::map<OpBase*, int> ComputeDepCounts(OpBase* op) {
    std::map<OpBase*, int> ret;

    std::deque<OpBase*> queue;
    queue.push_back(op);
    std::unordered_set<OpBase*> visited;
    visited.insert(op);
    while (!queue.empty()) {
      OpBase* candidate = queue.front();
      queue.pop_front();
X
Xin Pan 已提交
104
      for (auto it : candidate->pre_ops_) {
X
Xin Pan 已提交
105 106 107 108 109 110 111
        for (OpBase* pre_op : it.second) {
          if (!pre_op) continue;
          if (visited.find(pre_op) == visited.end()) {
            visited.insert(pre_op);
            queue.push_back(pre_op);
          }
          ret[pre_op] += 1;
112 113 114 115 116 117 118 119 120 121 122 123
        }
      }
    }
    return ret;
  }
};

framework::LoDTensor& VarBase::Grad() {
  VLOG(3) << "get var grad " << var_desc_->Name();
  return *grads_->GetMutable<framework::LoDTensor>();
}

X
Xin Pan 已提交
124 125
std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
  if (!grad_op_desc_) {
126
    LOG(WARNING) << "op with no grad: " << op_desc_->Type();
X
Xin Pan 已提交
127
    return {};
128 129 130
  }
  VLOG(3) << "op grad " << grad_op_desc_->Type();

X
Xin Pan 已提交
131
  std::vector<std::unique_ptr<framework::Variable>> tmp_vars;
X
Xin Pan 已提交
132 133 134 135
  std::map<std::string, std::vector<framework::Variable*>> grad_outputs;
  for (auto it : grad_output_vars_) {
    auto& outputs = grad_outputs[it.first];
    for (size_t i = 0; i < it.second.size(); ++i) {
136 137 138 139 140 141
      // Allocate a new variable
      Variable* tmp_var = new framework::Variable();
      tmp_var->GetMutable<framework::LoDTensor>();

      tmp_vars.emplace_back(tmp_var);
      outputs.push_back(tmp_var);
142 143 144
    }
  }

X
Xin Pan 已提交
145
  framework::RuntimeContext ctx(grad_input_vars_, grad_outputs);
146

147
  // No need to do compile time infer shape here.
X
Xin Pan 已提交
148
  // grad_op_desc_->InferShape(*block_);
149
  grad_op_desc_->InferVarType(block_);
X
Xin Pan 已提交
150

151 152
  std::unique_ptr<framework::OperatorBase> opbase =
      framework::OpRegistry::CreateOp(*grad_op_desc_);
X
Xin Pan 已提交
153 154 155 156 157 158 159 160 161
  framework::OperatorWithKernel* op_kernel =
      dynamic_cast<framework::OperatorWithKernel*>(opbase.get());
  PADDLE_ENFORCE_NOT_NULL(op_kernel, "only support op with kernel");

  framework::Scope scope;
  platform::CPUPlace place;
  PreparedOp p = PreparedOp::Prepare(ctx, *op_kernel, place);
  p.op.RuntimeInferShape(scope, place, ctx);
  p.func(framework::ExecutionContext(p.op, scope, *p.dev_ctx, p.ctx));
X
Xin Pan 已提交
162 163 164 165

  for (auto it : grad_output_vars_) {
    auto& outputs = grad_outputs[it.first];
    auto& origin_outputs = it.second;
166

X
Xin Pan 已提交
167
    for (size_t i = 0; i < outputs.size(); ++i) {
M
minqiyang 已提交
168 169
      framework::Variable* orig_grad = origin_outputs[i];
      AddTo(outputs[i], orig_grad);
170 171
    }
  }
X
Xin Pan 已提交
172
  return input_vars_;
173 174
}

X
Xin Pan 已提交
175
void VarBase::RunBackward() {
176
  if (!pre_op_) return;
X
Xin Pan 已提交
177

X
Xin Pan 已提交
178 179 180 181
  auto grads_t = grads_->GetMutable<framework::LoDTensor>();
  float* data = grads_t->mutable_data<float>(platform::CPUPlace());
  std::fill(data, data + grads_t->numel(), 1.0);

X
Xin Pan 已提交
182 183 184
  PADDLE_ENFORCE(
      grads_ ==
      pre_op_->output_vars_[pre_op_out_name_][pre_op_out_idx_]->grads_);
X
Xin Pan 已提交
185
  Autograd().RunBackward(this);
186 187
}

X
Xin Pan 已提交
188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204
void PyLayer::RegisterFunc(int func_id, const py::object& py_func) {
  py_funcs_[func_id] = py_func;
}

std::vector<VarBase*> PyLayer::Apply(int func_id,
                                     const std::vector<VarBase>& inputs) {
  std::vector<framework::LoDTensor> tensor_inputs;
  std::vector<VarBase*> ret;

  for (const VarBase& in : inputs) {
    tensor_inputs.push_back(in.var_->Get<framework::LoDTensor>());
  }
  PADDLE_ENFORCE(py_funcs_.find(func_id) != py_funcs_.end());
  CallPythonFunc(py_funcs_[func_id], tensor_inputs, &ret);
  return ret;
}

205 206
}  // namespace imperative
}  // namespace paddle