layer.h 6.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
// 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.

#pragma once

X
Xin Pan 已提交
17 18 19 20 21 22 23
// clang-format off
#include "paddle/fluid/framework/python_headers.h"
// clang-format on

#include <map>     // NOLINT
#include <string>  // NOLINT
#include <vector>  // NOLINT
M
minqiyang 已提交
24

25 26 27 28 29
#include "paddle/fluid/framework/op_desc.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/framework/var_desc.h"
#include "paddle/fluid/platform/enforce.h"

M
minqiyang 已提交
30 31
#include "paddle/fluid/imperative/type_defs.h"

32 33 34
namespace paddle {
namespace imperative {

X
Xin Pan 已提交
35 36
namespace py = ::pybind11;

X
Xin Pan 已提交
37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
class PreparedOp {
 public:
  PreparedOp(const framework::OperatorBase& op,
             const framework::RuntimeContext& ctx,
             framework::OperatorWithKernel::OpKernelFunc func,
             platform::DeviceContext* dev_ctx)
      : op(op), ctx(ctx), func(func), dev_ctx(dev_ctx) {}

  static PreparedOp Prepare(const framework::RuntimeContext& ctx,
                            const framework::OperatorWithKernel& op,
                            const platform::Place& place) {
    platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
    auto* dev_ctx = pool.Get(place);

    // check if op[type] has kernel registered.
    auto& all_op_kernels = op.AllOpKernels();
    auto kernels_iter = all_op_kernels.find(op.Type());
    if (kernels_iter == all_op_kernels.end()) {
      PADDLE_THROW(
          "There are no kernels which are registered in the %s operator.",
          op.Type());
    }

    framework::OperatorWithKernel::OpKernelMap& kernels = kernels_iter->second;

    auto expected_kernel_key = op.GetExpectedKernelType(
X
Xin Pan 已提交
63
        framework::ExecutionContext(op, framework::Scope(), *dev_ctx, ctx));
X
Xin Pan 已提交
64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88
    VLOG(3) << "expected_kernel_key:" << expected_kernel_key;

    auto kernel_iter = kernels.find(expected_kernel_key);
#ifdef PADDLE_WITH_MKLDNN
    // workaround for missing MKLDNN kernel when FLAGS_use_mkldnn env var is set
    if (kernel_iter == kernels.end() &&
        expected_kernel_key.library_type_ == framework::LibraryType::kMKLDNN) {
      VLOG(3) << "missing MKLDNN kernel: fallbacking to PLAIN one";
      expected_kernel_key.library_type_ = framework::LibraryType::kPlain;
      expected_kernel_key.data_layout_ = framework::DataLayout::kAnyLayout;
      kernel_iter = kernels.find(expected_kernel_key);
    }
#endif
    if (kernel_iter == kernels.end()) {
      PADDLE_THROW("op %s does not have kernel for %s", op.Type(),
                   KernelTypeToString(expected_kernel_key));
    }
    return PreparedOp(op, ctx, kernel_iter->second, dev_ctx);
  }

  const framework::OperatorBase& op;
  const framework::RuntimeContext& ctx;
  framework::OperatorWithKernel::OpKernelFunc func;
  platform::DeviceContext* dev_ctx;
};
X
polish  
Xin Pan 已提交
89

90 91
class OpBase;

M
minqiyang 已提交
92 93 94 95 96
/* The wrapper for Variable which holds a Variable and a VarBase of its
 * gradient. This object should be managed totally by Python intepreter.
 *
 * Nearly all interface should be implemented in C++.
 */
97 98
class VarBase {
 public:
99
  VarBase() : VarBase(new framework::Variable(), new VarBase(true)) {}
X
polish  
Xin Pan 已提交
100 101

  // Owns `var` and `grad`
102
  VarBase(framework::Variable* var, VarBase* grad)
M
minqiyang 已提交
103
      : pre_op_(nullptr),
104
        pre_op_out_name_(),
M
minqiyang 已提交
105 106
        pre_op_out_idx_(-1),
        var_desc_(nullptr),
X
polish  
Xin Pan 已提交
107 108
        var_(var),
        grads_(grad),
M
minqiyang 已提交
109 110 111
        stop_gradient_(false) {}

  explicit VarBase(bool stop_gradient)
112
      : pre_op_(nullptr),
113
        pre_op_out_name_(),
114 115
        pre_op_out_idx_(-1),
        var_desc_(nullptr),
X
Xin Pan 已提交
116
        var_(new framework::Variable()),
M
minqiyang 已提交
117
        grads_(stop_gradient ? nullptr : new VarBase(true)),
118
        stop_gradient_(stop_gradient) {}
119

M
minqiyang 已提交
120 121 122 123 124 125 126 127 128
  virtual ~VarBase() {
    if (var_) {
      delete var_;
    }

    if (grads_) {
      delete grads_;
    }
  }
129

X
Xin Pan 已提交
130
  void RunBackward();
131

M
minqiyang 已提交
132
  framework::LoDTensor& GradValue();
133

M
minqiyang 已提交
134 135 136 137 138 139 140
  inline std::string GradName() const {
    PADDLE_ENFORCE(
        var_desc_,
        "Couldn't get gradient variable's name, please call backward() first");
    return string::Sprintf("%s@IGrad", var_desc_->Name());
  }

141
  OpBase* pre_op_;
X
Xin Pan 已提交
142
  std::string pre_op_out_name_;
143 144 145
  int pre_op_out_idx_;

  framework::VarDesc* var_desc_;
M
minqiyang 已提交
146

M
minqiyang 已提交
147 148
  framework::Variable* var_;
  VarBase* grads_;
149 150

  bool stop_gradient_;
151 152
};

M
minqiyang 已提交
153 154 155
/* The wrapper for OpDesc which holds a OpDesc and a OpDesc of its
 * gradient. This object should be managed totally by Python intepreter.
 */
156 157
class OpBase {
 public:
X
Xin Pan 已提交
158 159 160
  OpBase()
      : op_desc_(nullptr),
        forward_id_(-1),
X
polish  
Xin Pan 已提交
161
        grad_op_desc_(nullptr),
X
Xin Pan 已提交
162
        backward_id_(-1) {}
163 164 165 166 167

  virtual ~OpBase() {
    if (grad_op_desc_) delete grad_op_desc_;
  }

X
Xin Pan 已提交
168
  std::map<std::string, std::vector<VarBase*>> ApplyGrad();
169

X
polish  
Xin Pan 已提交
170 171
  // One of `op_desc_` or `forward_id_` is set, not both.
  // For pure python PyLayer, use `forward_id_`, otherwise, use op_desc_.
172
  framework::OpDesc* op_desc_;
X
Xin Pan 已提交
173
  int forward_id_;
X
polish  
Xin Pan 已提交
174 175
  // When has backward, one of `grad_op_desc_` or `backward_id_` is set,
  // not both.
176
  framework::OpDesc* grad_op_desc_;
X
Xin Pan 已提交
177
  int backward_id_;
X
Xin Pan 已提交
178

M
minqiyang 已提交
179 180 181
  VarBasePtrMap input_vars_;
  VarBasePtrMap output_vars_;
  OpBasePtrMap pre_ops_;
X
Xin Pan 已提交
182
  std::map<std::string, std::vector<int>> pre_ops_out_idx_;
183

M
minqiyang 已提交
184 185
  framework::VariableValueMap grad_input_vars_;
  framework::VariableValueMap grad_output_vars_;
186 187 188 189 190 191 192 193 194 195 196
  framework::BlockDesc* block_;
};

class Layer {
 public:
  virtual ~Layer() {}

  virtual std::vector<VarBase> Forward(const std::vector<VarBase>& inputs) {
    std::vector<VarBase> vars;
    return vars;
  }
X
Xin Pan 已提交
197
};
198

X
Xin Pan 已提交
199 200 201 202
class PyLayer {
 public:
  virtual ~PyLayer() {}

X
Xin Pan 已提交
203 204 205
  static constexpr char* kFwdInp = "X";
  static constexpr char* kFwdOut = "Out";

X
Xin Pan 已提交
206
  static void RegisterFunc(int func_id, const py::object& py_func);
X
Xin Pan 已提交
207

X
polish  
Xin Pan 已提交
208 209
  static int NumFuncs();

X
Xin Pan 已提交
210
  static std::vector<VarBase*> Apply(int func_id,
X
Xin Pan 已提交
211 212
                                     const std::vector<VarBase*>& inputs);

X
polish  
Xin Pan 已提交
213 214
  static std::vector<framework::Variable*> ApplyGrad(
      int func_id, const std::vector<framework::Variable*>& inputs);
215

X
polish  
Xin Pan 已提交
216 217 218
 private:
  static std::vector<framework::Variable*> CallPythonFunc(
      const py::object& callable, const std::vector<framework::Variable*>& ins);
219 220 221 222
};

}  // namespace imperative
}  // namespace paddle