backward.cc 8.5 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->AddOp(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;
D
dongzhihong 已提交
127
      auto& dup_op = dup_output_op.second;
Y
Yan Chunwei 已提交
128
      // no duplicate output
D
dongzhihong 已提交
129 130
      if (dup_op.size() == 1) continue;

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

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

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

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

        // If part of input gradient of that operator is not calculated, fill
        // zero variables to that input gradient.
Y
Yu Yang 已提交
166 167
        net->AddOp(OpRegistry::CreateOp("fill_zeros_like", {{"Src", {prefix}}},
                                        {{"Dst", {grad_input}}}, {}));
168
      }
Y
Yu Yang 已提交
169 170 171
      return false;
    });

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

Y
Yan Chunwei 已提交
180 181 182 183 184 185 186 187 188 189 190 191
    // process recurrent gradient op as a special operator.
    if (forwardOp.Type() == "recurrent_op") {
      // 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 已提交
192
          BackwardRecursive(stepnet_op, no_grad_names, uniq_id));
Y
Yan Chunwei 已提交
193 194
    }

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

Y
Yu Yang 已提交
206
// See header for comments
Y
Yu Yang 已提交
207
std::unique_ptr<OperatorBase> Backward(
Y
Yu Yang 已提交
208
    const OperatorBase& forwardOp,
Y
Yu Yang 已提交
209 210 211 212
    const std::unordered_set<std::string>& no_grad_vars) {
  std::unordered_set<std::string> no_grad_names;
  no_grad_names.reserve(no_grad_vars.size());

213
  no_grad_names.insert(std::string(kEmptyVarName) + kGradVarSuffix);
214

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

Y
Yu Yang 已提交
222 223
}  // namespace framework
}  // namespace paddle