backward.cc 7.1 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>
18
#include "paddle/framework/op_registry.h"
Y
Yan Chunwei 已提交
19
#include "paddle/operators/net_op.h"
Y
Yu Yang 已提交
20 21 22 23

namespace paddle {
namespace framework {

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

static bool AllInSet(
Y
Yu Yang 已提交
34
    const std::map<std::string, std::vector<std::string>>& names,
Y
Yu Yang 已提交
35
    const std::string& suffix, const std::unordered_set<std::string>& set) {
36 37 38 39
  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 已提交
40
  });
41
  return all_in_set;
Y
Yu Yang 已提交
42 43
}

Y
Yu Yang 已提交
44
static std::shared_ptr<OperatorBase> NOP() {
Y
Yan Chunwei 已提交
45
  auto net_op = std::make_shared<operators::NetOp>();
Q
qiaolongfei 已提交
46
  net_op->SetType("@NOP@");
Y
Yu Yang 已提交
47 48 49 50
  net_op->CompleteAddOp();
  return net_op;
}

Y
Yu Yang 已提交
51 52 53 54
//  Get backward operator from a forward operator, recursively implementation.
//
//  no_grad_names the gradient variable names without gradient calculating.
//
55 56 57
//  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 已提交
58 59 60 61 62 63 64 65 66
//
//  returns The backward operator. For simple situation, it is a simple
//  operator. For complex situation, it is a NetOp.
//
//  See Backward.h for details
static std::shared_ptr<OperatorBase> BackwardRecursive(
    const OperatorBase& forwardOp,
    std::unordered_set<std::string>& no_grad_names, size_t& uniq_id);
std::shared_ptr<OperatorBase> BackwardRecursive(
Y
Yu Yang 已提交
67 68
    const OperatorBase& forwardOp,
    std::unordered_set<std::string>& no_grad_names, size_t& uniq_id) {
Y
Yu Yang 已提交
69 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
  //  too much time for calculation, but it is useful for simplifying logic.
Q
qiaolongfei 已提交
72
  if (AllInSet(forwardOp.Inputs(), kGradVarSuffix, no_grad_names)) {
Y
Yu Yang 已提交
73
    return NOP();
Y
Yu Yang 已提交
74 75
  }

76 77
  //  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 已提交
78
  //  `no_grad_names` set. Return an NOP.
Q
qiaolongfei 已提交
79 80
  if (AllInSet(forwardOp.Outputs(), kGradVarSuffix, no_grad_names)) {
    ForEachVarName(forwardOp.Inputs(),
Y
Yu Yang 已提交
81 82 83 84
                   [&no_grad_names](const std::string& name) -> bool {
                     no_grad_names.insert(GradVarName(name));
                     return false;
                   });
Y
Yu Yang 已提交
85
    return NOP();
Y
Yu Yang 已提交
86 87
  }

Y
Yu Yang 已提交
88
  // Returned gradient network
Y
Yan Chunwei 已提交
89
  auto net = std::make_shared<operators::NetOp>();
Y
Yu Yang 已提交
90 91

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

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

    size_t local_op_id = 0;
Y
Yu Yang 已提交
100
    // reversely travel forwardNet
Y
Yu Yang 已提交
101
    for (auto it = forwardNet.ops_.rbegin(); it != forwardNet.ops_.rend();
Y
Yu Yang 已提交
102
         ++it, ++local_op_id) {
D
dongzhihong 已提交
103
      auto fwd = *it;
Y
Yu Yang 已提交
104
      auto bwd = BackwardRecursive(*fwd, no_grad_names, uniq_id);
D
dongzhihong 已提交
105
      net->AddOp(bwd);
Q
qiaolongfei 已提交
106
      ForEachVarName(bwd->Outputs(),
Y
Yu Yang 已提交
107 108 109 110
                     [&dup_output_ops, local_op_id](const std::string& out) {
                       dup_output_ops[out].emplace_back(local_op_id);
                       return false;
                     });
D
dongzhihong 已提交
111
    }
Y
Yu Yang 已提交
112
    // Get unique ID for this method.
D
dongzhihong 已提交
113
    auto uid = uniq_id++;
D
dongzhihong 已提交
114
    // TODO(dzh): more comment
Y
Yu Yang 已提交
115 116
    using Pos = std::pair<size_t, std::shared_ptr<OperatorBase>>;
    std::list<Pos> insert_position;
D
dongzhihong 已提交
117
    for (auto& dup_output_op : dup_output_ops) {
D
dongzhihong 已提交
118
      const std::string& name = dup_output_op.first;
D
dongzhihong 已提交
119 120 121 122 123 124
      auto& dup_op = dup_output_op.second;
      if (dup_op.size() == 1) continue;
      std::vector<std::string> dup_outputs;

      for (size_t i = 0; i < dup_op.size(); ++i) {
        auto op_offset = dup_op[i];
D
dongzhihong 已提交
125 126 127
        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 已提交
128
      }
Y
Yu Yang 已提交
129
      insert_position.push_back(
D
dongzhihong 已提交
130 131
          {dup_op.back(),
           OpRegistry::CreateOp(
Y
Yu Yang 已提交
132
               "add", {{"X", {dup_outputs}}}, {{"Out", {name}}},
D
dongzhihong 已提交
133
               {{"input_format",
Y
Yu Yang 已提交
134
                 std::vector<int>{0, static_cast<int>(dup_outputs.size())}}})});
D
dongzhihong 已提交
135
    }
Y
Yu Yang 已提交
136 137

    insert_position.sort(
D
dongzhihong 已提交
138
        [](const Pos& l, const Pos& r) { return l.first > r.first; });
Y
Yu Yang 已提交
139 140

    for (auto& pos : insert_position) {
Y
Yu Yang 已提交
141
      net->InsertOp(pos.first + 1, pos.second);
D
dongzhihong 已提交
142 143
    }

Y
Yu Yang 已提交
144
  } else {
145
    std::shared_ptr<OperatorBase> grad_op = OpRegistry::CreateGradOp(forwardOp);
Y
Yu Yang 已提交
146

Q
qiaolongfei 已提交
147 148
    ForEachVarName(grad_op->Inputs(), [&no_grad_names, &net,
                                       grad_op](const std::string& grad_input) {
149
      if (no_grad_names.count(grad_input)) {
Y
Yu Yang 已提交
150
        // +1 for \0
151
        std::string prefix = grad_input.substr(
Y
Yu Yang 已提交
152
            0, grad_input.size() - sizeof(kGradVarSuffix) / sizeof(char) + 1);
Q
qiaolongfei 已提交
153
        grad_op->Rename(grad_input, prefix + kZeroVarSuffix);
Y
Yu Yang 已提交
154 155 156

        // If part of input gradient of that operator is not calculated, fill
        // zero variables to that input gradient.
Y
Yu Yang 已提交
157 158
        net->AddOp(OpRegistry::CreateOp("fill_zeros_like", {{"Src", {prefix}}},
                                        {{"Dst", {grad_input}}}, {}));
159
      }
Y
Yu Yang 已提交
160 161 162
      return false;
    });

Q
qiaolongfei 已提交
163 164
    ForEachVarName(grad_op->Outputs(),
                   [&no_grad_names, &grad_op](const std::string& grad_output) {
Y
Yu Yang 已提交
165
                     if (no_grad_names.count(grad_output)) {
Q
qiaolongfei 已提交
166
                       grad_op->Rename(grad_output, kEmptyVarName);
Y
Yu Yang 已提交
167 168 169
                     }
                     return false;
                   });
Y
Yu Yang 已提交
170 171 172 173

    if (net->ops_.empty()) {  // Current no aux op is added to network
      return grad_op;
    }
F
fengjiayi 已提交
174
    net->AddOp(grad_op);
Y
Yu Yang 已提交
175
  }
Q
qiaolongfei 已提交
176
  net->SetType("@GENERATED_BACKWARD@");
Y
Yu Yang 已提交
177
  net->CompleteAddOp();
Y
Yu Yang 已提交
178
  return net;
Y
Yu Yang 已提交
179 180
}

Y
Yu Yang 已提交
181 182
// See header for comments
std::shared_ptr<OperatorBase> Backward(
Y
Yu Yang 已提交
183
    const OperatorBase& forwardOp,
Y
Yu Yang 已提交
184 185 186 187
    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());

188
  no_grad_names.insert(std::string(kEmptyVarName) + kGradVarSuffix);
189

Y
Yu Yang 已提交
190
  for (auto& name : no_grad_vars) {
191
    no_grad_names.insert(name + kGradVarSuffix);
Y
Yu Yang 已提交
192
  }
Y
Yu Yang 已提交
193
  size_t uid = 0;
Y
Yu Yang 已提交
194
  return BackwardRecursive(forwardOp, no_grad_names, uid);
Y
Yu Yang 已提交
195
}
Y
Yi Wang 已提交
196

Y
Yu Yang 已提交
197 198
}  // namespace framework
}  // namespace paddle