backward.cc 8.6 KB
Newer Older
Y
Yu Yang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* 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. */

15
#include "paddle/framework/backward.h"
D
dongzhihong 已提交
16

D
dongzhihong 已提交
17
#include <list>
Y
Yu Yang 已提交
18 19
#include <memory>

20
#include "paddle/framework/op_registry.h"
Y
Yan Chunwei 已提交
21
#include "paddle/operators/net_op.h"
Y
Yan Chunwei 已提交
22
#include "paddle/operators/recurrent_op.h"
Y
Yu Yang 已提交
23 24 25 26

namespace paddle {
namespace framework {

Y
Yu Yang 已提交
27
template <typename Map, typename T>
Q
qiaolongfei 已提交
28
static void ForEachVarName(const Map& names, T callback) {
Y
Yu Yang 已提交
29
  for (auto& name : names) {
Y
Yu Yang 已提交
30
    for (auto& n : name.second) {
31
      if (callback(n)) return;
Y
Yu Yang 已提交
32 33
    }
  }
Y
Yu Yang 已提交
34 35
}

Y
Yan Chunwei 已提交
36
// return whether all the names + suffixes in the set
Y
Yu Yang 已提交
37
static bool AllInSet(
Y
Yu Yang 已提交
38
    const std::map<std::string, std::vector<std::string>>& names,
Y
Yu Yang 已提交
39
    const std::string& suffix, const std::unordered_set<std::string>& set) {
40 41 42 43
  bool all_in_set = true;
  ForEachVarName(names, [&all_in_set, &set, &suffix](const std::string& n) {
    all_in_set = set.find(n + suffix) != set.end();
    return !all_in_set;
Y
Yu Yang 已提交
44
  });
45
  return all_in_set;
Y
Yu Yang 已提交
46 47
}

Y
Yu Yang 已提交
48 49
static std::unique_ptr<OperatorBase> NOP() {
  auto net_op = new operators::NetOp();
Q
qiaolongfei 已提交
50
  net_op->SetType("@NOP@");
Y
Yu Yang 已提交
51
  net_op->CompleteAddOp();
Y
Yu Yang 已提交
52
  return std::unique_ptr<OperatorBase>(net_op);
Y
Yu Yang 已提交
53 54
}

Y
Yan Chunwei 已提交
55
//  Get backward operator from a forward operator, a recursive implementation.
Y
Yu Yang 已提交
56 57 58
//
//  no_grad_names the gradient variable names without gradient calculating.
//
59 60 61
//  uniq_id is a unique index used inside recursively calling
//  BackwardRecursive. use `uid = uniq_id++;` to get the unique index, and
//  pass `uniq_id` through recursive calling.
Y
Yu Yang 已提交
62
//
Y
Yan Chunwei 已提交
63 64
//  returns The backward operator. In a simple situation, it may be a simple
//  operator, in a complex situation, it maybe a NetOp.
Y
Yu Yang 已提交
65 66
//
//  See Backward.h for details
Y
Yu Yang 已提交
67
static std::unique_ptr<OperatorBase> BackwardRecursive(
Y
Yu Yang 已提交
68 69
    const OperatorBase& forwardOp,
    std::unordered_set<std::string>& no_grad_names, size_t& uniq_id) {
Y
Yu Yang 已提交
70 71
  //  If all input gradients of forwarding operator do not need to calculate,
  //  just return an NOP. Not return null ptr because NOP does not take
Q
typo  
qiaolongfei 已提交
72
  //  too much time for calculation, but it is useful for simplifying logic.
73
  if (AllInSet(forwardOp.Inputs() /*names*/, kGradVarSuffix /*suffix*/,
Y
Yan Chunwei 已提交
74
               no_grad_names /*set*/)) {
Y
Yu Yang 已提交
75
    return NOP();
Y
Yu Yang 已提交
76 77
  }

78 79
  //  All output gradients of forwarding operator do not need to calculate.
  //  Then all input gradients cannot be computed at all, and we put them into
Y
Yu Yang 已提交
80
  //  `no_grad_names` set. Return an NOP.
Q
qiaolongfei 已提交
81
  if (AllInSet(forwardOp.Outputs() /*names*/, kGradVarSuffix /*suffix*/,
Y
Yan Chunwei 已提交
82
               no_grad_names /*set*/)) {
Q
qiaolongfei 已提交
83
    ForEachVarName(forwardOp.Inputs(),
Y
Yu Yang 已提交
84 85 86 87
                   [&no_grad_names](const std::string& name) -> bool {
                     no_grad_names.insert(GradVarName(name));
                     return false;
                   });
Y
Yu Yang 已提交
88
    return NOP();
Y
Yu Yang 已提交
89 90
  }

Y
Yu Yang 已提交
91
  // Returned gradient network
Y
Yu Yang 已提交
92
  auto net = std::unique_ptr<operators::NetOp>(new operators::NetOp());
Y
Yu Yang 已提交
93 94

  if (forwardOp.IsNetOp()) {
Y
Yu Yang 已提交
95
    // Because forwardOp is a net op, it can static_cast.
Y
Yan Chunwei 已提交
96
    auto& forwardNet = static_cast<const operators::NetOp&>(forwardOp);
Y
Yu Yang 已提交
97

98
    // Map from output gradient variable name to operator's indices in
Y
Yan Chunwei 已提交
99
    // backward net's ops_. That operator generates that variable.
Y
Yu Yang 已提交
100 101 102
    std::unordered_map<std::string, std::vector<size_t>> dup_output_ops;

    size_t local_op_id = 0;
Y
Yan Chunwei 已提交
103
    // reversely travel forwardNet and collect all duplicate outputs.
Y
Yu Yang 已提交
104
    for (auto it = forwardNet.ops_.rbegin(); it != forwardNet.ops_.rend();
Y
Yu Yang 已提交
105
         ++it, ++local_op_id) {
Y
Yu Yang 已提交
106
      auto& fwd = *it;
Y
Yu Yang 已提交
107
      auto bwd = BackwardRecursive(*fwd, no_grad_names, uniq_id);
Q
qiaolongfei 已提交
108
      ForEachVarName(bwd->Outputs(),
Y
Yu Yang 已提交
109 110 111 112
                     [&dup_output_ops, local_op_id](const std::string& out) {
                       dup_output_ops[out].emplace_back(local_op_id);
                       return false;
                     });
Y
Yu Yang 已提交
113
      net->AppendOp(std::move(bwd));
D
dongzhihong 已提交
114
    }
Y
Yu Yang 已提交
115
    // Get unique ID for this method.
D
dongzhihong 已提交
116
    auto uid = uniq_id++;
D
dongzhihong 已提交
117
    // TODO(dzh): more comment
Y
Yan Chunwei 已提交
118 119 120 121 122
    // multiple operators which have the same output (y for example) may
    // overwrite the same y variable when backward, special operations are token
    // to handle this case. For each duplicate output, rename it to an alias
    // (original name with a offset), append an `add` op for its operator,
    // and finally sum all the alias variable to the final output variable y.
Y
Yu Yang 已提交
123
    using Pos = std::pair<size_t, std::unique_ptr<OperatorBase>>;
Y
Yu Yang 已提交
124
    std::list<Pos> insert_position;
D
dongzhihong 已提交
125
    for (auto& dup_output_op : dup_output_ops) {
D
dongzhihong 已提交
126
      const std::string& name = dup_output_op.first;
Q
qijun 已提交
127 128 129
      // duplicate @Empty@ don't need to be added
      if (name == kEmptyVarName) continue;

D
dongzhihong 已提交
130
      auto& dup_op = dup_output_op.second;
Y
Yan Chunwei 已提交
131
      // no duplicate output
D
dongzhihong 已提交
132 133
      if (dup_op.size() == 1) continue;

Y
Yan Chunwei 已提交
134 135
      // process the duplicate outputs
      std::vector<std::string> dup_outputs;
D
dongzhihong 已提交
136
      for (size_t i = 0; i < dup_op.size(); ++i) {
Y
Yan Chunwei 已提交
137
        // rename each duplicate output to an alias
D
dongzhihong 已提交
138
        auto op_offset = dup_op[i];
D
dongzhihong 已提交
139 140 141
        dup_outputs.push_back(name + "@RENAME@" + std::to_string(uid) + "@" +
                              std::to_string(i));
        net->ops_[op_offset]->Rename(name, dup_outputs.back());
D
dongzhihong 已提交
142
      }
Y
Yan Chunwei 已提交
143
      // collect all the offset to append `add` op for each alias
Y
Yu Yang 已提交
144
      insert_position.push_back(
Y
Yu Yang 已提交
145 146
          {dup_op.back(), OpRegistry::CreateOp("add", {{"X", {dup_outputs}}},
                                               {{"Out", {name}}}, {})});
D
dongzhihong 已提交
147
    }
Y
Yu Yang 已提交
148

Y
Yan Chunwei 已提交
149
    // make sure the inserted `add` ops follow the BFS order.
Y
Yu Yang 已提交
150
    insert_position.sort(
D
dongzhihong 已提交
151
        [](const Pos& l, const Pos& r) { return l.first > r.first; });
Y
Yu Yang 已提交
152 153

    for (auto& pos : insert_position) {
Y
Yu Yang 已提交
154
      net->InsertOp(pos.first + 1, std::move(pos.second));
D
dongzhihong 已提交
155
    }
Y
Yu Yang 已提交
156
  } else {
Y
Yu Yang 已提交
157
    std::unique_ptr<OperatorBase> grad_op(OpRegistry::CreateGradOp(forwardOp));
Y
Yu Yang 已提交
158

Y
Yu Yang 已提交
159 160
    ForEachVarName(grad_op->Inputs(), [&no_grad_names, &net, &grad_op](
                                          const std::string& grad_input) {
161
      if (no_grad_names.count(grad_input)) {
Y
Yu Yang 已提交
162
        // +1 for \0
163
        std::string prefix = grad_input.substr(
Y
Yu Yang 已提交
164
            0, grad_input.size() - sizeof(kGradVarSuffix) / sizeof(char) + 1);
Q
qiaolongfei 已提交
165
        grad_op->Rename(grad_input, prefix + kZeroVarSuffix);
Y
Yu Yang 已提交
166 167 168

        // If part of input gradient of that operator is not calculated, fill
        // zero variables to that input gradient.
D
dangqingqing 已提交
169 170
        net->AppendOp(OpRegistry::CreateOp("fill_zeros_like", {{"X", {prefix}}},
                                           {{"Y", {grad_input}}}, {}));
171
      }
Y
Yu Yang 已提交
172 173 174
      return false;
    });

Q
qiaolongfei 已提交
175 176
    ForEachVarName(grad_op->Outputs(),
                   [&no_grad_names, &grad_op](const std::string& grad_output) {
Y
Yu Yang 已提交
177
                     if (no_grad_names.count(grad_output)) {
Q
qiaolongfei 已提交
178
                       grad_op->Rename(grad_output, kEmptyVarName);
Y
Yu Yang 已提交
179 180 181
                     }
                     return false;
                   });
Y
Yu Yang 已提交
182

Y
Yan Chunwei 已提交
183
    // process recurrent gradient op as a special operator.
184
    if (forwardOp.Type() == "recurrent") {
Y
Yan Chunwei 已提交
185 186 187 188 189 190 191 192 193 194
      // NOTE clean up cycle call somewhere (RNN's stepnet constains itself), or
      // this will result in infinite loop.
      const auto& rnnop =
          *static_cast<const operators::RecurrentOp*>(&forwardOp);
      auto rnn_grad_op =
          static_cast<operators::RecurrentGradientOp*>(grad_op.get());
      const auto& stepnet_op =
          *static_cast<const OperatorBase*>(&rnnop.stepnet());
      // create stepnet's gradient op
      rnn_grad_op->set_stepnet(
Y
Yu Yang 已提交
195
          BackwardRecursive(stepnet_op, no_grad_names, uniq_id));
Y
Yan Chunwei 已提交
196 197
    }

Y
Yu Yang 已提交
198 199 200
    if (net->ops_.empty()) {  // Current no aux op is added to network
      return grad_op;
    }
Y
Yu Yang 已提交
201
    net->AppendOp(std::move(grad_op));
Y
Yu Yang 已提交
202
  }
Q
qiaolongfei 已提交
203
  net->SetType("@GENERATED_BACKWARD@");
Y
Yu Yang 已提交
204
  net->CompleteAddOp();
Y
Yu Yang 已提交
205 206 207
  return std::unique_ptr<OperatorBase>(
      static_cast<OperatorBase*>(net.release()));
}
Y
Yu Yang 已提交
208

Y
Yu Yang 已提交
209
// See header for comments
Y
Yu Yang 已提交
210
std::unique_ptr<OperatorBase> Backward(
Y
Yu Yang 已提交
211
    const OperatorBase& forwardOp,
Y
Yu Yang 已提交
212 213
    const std::unordered_set<std::string>& no_grad_vars) {
  std::unordered_set<std::string> no_grad_names;
Q
qijun 已提交
214
  no_grad_names.reserve(no_grad_vars.size() + 1);
Y
Yu Yang 已提交
215

216
  no_grad_names.insert(std::string(kEmptyVarName) + kGradVarSuffix);
217

Y
Yu Yang 已提交
218
  for (auto& name : no_grad_vars) {
219
    no_grad_names.insert(name + kGradVarSuffix);
Y
Yu Yang 已提交
220
  }
Y
Yu Yang 已提交
221
  size_t uid = 0;
Y
Yu Yang 已提交
222
  return BackwardRecursive(forwardOp, no_grad_names, uid);
Y
Yu Yang 已提交
223
}
Y
Yi Wang 已提交
224

Y
Yu Yang 已提交
225 226
}  // namespace framework
}  // namespace paddle