layer.h 5.3 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
#include <map>
18 19
#include <string>
#include <vector>
M
minqiyang 已提交
20

21 22 23 24 25
#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 已提交
26 27
#include "paddle/fluid/imperative/type_defs.h"

28 29 30
namespace paddle {
namespace imperative {

X
Xin Pan 已提交
31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
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 已提交
57
        framework::ExecutionContext(op, framework::Scope(), *dev_ctx, ctx));
X
Xin Pan 已提交
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82
    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;
};
83 84
class OpBase;

M
minqiyang 已提交
85 86 87 88 89
/* 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++.
 */
90 91
class VarBase {
 public:
M
minqiyang 已提交
92 93 94 95 96
  VarBase()
      : pre_op_(nullptr),
        pre_op_out_idx_(-1),
        var_desc_(nullptr),
        var_(new framework::Variable()),
M
minqiyang 已提交
97
        grads_(new VarBase(true)),
M
minqiyang 已提交
98 99 100
        stop_gradient_(false) {}

  explicit VarBase(bool stop_gradient)
101 102 103
      : pre_op_(nullptr),
        pre_op_out_idx_(-1),
        var_desc_(nullptr),
X
Xin Pan 已提交
104
        var_(new framework::Variable()),
M
minqiyang 已提交
105
        grads_(stop_gradient ? nullptr : new VarBase(true)),
106
        stop_gradient_(stop_gradient) {}
107

M
minqiyang 已提交
108
  virtual ~VarBase() {}
109

X
Xin Pan 已提交
110
  void RunBackward();
111

M
minqiyang 已提交
112
  framework::LoDTensor& GradValue();
113

M
minqiyang 已提交
114 115 116 117 118 119 120
  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());
  }

121
  OpBase* pre_op_;
X
Xin Pan 已提交
122
  std::string pre_op_out_name_;
123 124 125
  int pre_op_out_idx_;

  framework::VarDesc* var_desc_;
M
minqiyang 已提交
126 127 128

  std::shared_ptr<framework::Variable> var_;
  std::shared_ptr<VarBase> grads_;
129 130

  bool stop_gradient_;
131 132
};

M
minqiyang 已提交
133 134 135
/* The wrapper for OpDesc which holds a OpDesc and a OpDesc of its
 * gradient. This object should be managed totally by Python intepreter.
 */
136 137
class OpBase {
 public:
X
Xin Pan 已提交
138
  OpBase() : op_desc_(nullptr), grad_op_desc_(nullptr) {}
139 140 141 142 143

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

X
Xin Pan 已提交
144
  std::map<std::string, std::vector<VarBase*>> ApplyGrad();
145 146 147

  framework::OpDesc* op_desc_;
  framework::OpDesc* grad_op_desc_;
X
Xin Pan 已提交
148

M
minqiyang 已提交
149 150 151
  VarBasePtrMap input_vars_;
  VarBasePtrMap output_vars_;
  OpBasePtrMap pre_ops_;
X
Xin Pan 已提交
152
  std::map<std::string, std::vector<int>> pre_ops_out_idx_;
153

M
minqiyang 已提交
154 155
  framework::VariableValueMap grad_input_vars_;
  framework::VariableValueMap grad_output_vars_;
156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
  framework::BlockDesc* block_;
};

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

  virtual std::vector<VarBase> Forward(const std::vector<VarBase>& inputs) {
    std::vector<VarBase> vars;
    return vars;
  }

  virtual void Backward() { LOG(ERROR) << "To support customize"; }
};

}  // namespace imperative
}  // namespace paddle