cond_op.cc 8.2 KB
Newer Older
Z
cond op  
zchen0211 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* 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. */

#include "paddle/operators/cond_op.h"
Z
zchen0211 已提交
16 17 18 19

#include <cstring>
#include <sstream>

Z
cond op  
zchen0211 已提交
20
#include "paddle/framework/op_registry.h"
Z
zchen0211 已提交
21
#include "paddle/operators/gather.h"
Z
cond op  
zchen0211 已提交
22
#include "paddle/operators/net_op.h"
Z
zchen0211 已提交
23
#include "paddle/operators/scatter.h"
Z
cond op  
zchen0211 已提交
24 25 26 27

namespace paddle {
namespace operators {

Z
zchen0211 已提交
28 29 30
using Scope = framework::Scope;
using Variable = framework::Variable;
using Tensor = framework::Tensor;
Z
zchen0211 已提交
31
using LoDTensor = framework::LoDTensor;
Z
zchen0211 已提交
32 33 34 35
using DDim = framework::DDim;

void CondOp::CreateScope(const Scope& scope) const {
  auto sub_scopes_var = scope.FindVar("SubScopes");
36 37
  PADDLE_ENFORCE_NOT_NULL(sub_scopes_var,
                          "Output(SubScopes) of CondOp should not be null.");
Z
zchen0211 已提交
38 39 40 41 42 43 44
  auto sub_scopes = sub_scopes_var->GetMutable<std::vector<Scope*>>();
  auto& sub_scope = scope.NewScope();
  sub_scopes->push_back(&sub_scope);
}

void CondOp::CreateIndexTensor(const Scope& scope) const {
  auto index_tensors_var = scope.FindVar("IndexTensors");
45 46
  PADDLE_ENFORCE_NOT_NULL(index_tensors_var,
                          "Output(IndexTensors) of CondOp should not be null.");
Z
zchen0211 已提交
47 48 49
  auto& index_tensors =
      *index_tensors_var->GetMutable<std::vector<LoDTensor>>();
  index_tensors.push_back(LoDTensor());
Z
zchen0211 已提交
50 51 52 53
}

void CondOp::InferShape(const Scope& scope) const {
  auto sub_scopes_var = scope.FindVar("SubScopes");
54 55
  PADDLE_ENFORCE_NOT_NULL(sub_scopes_var,
                          "Output(SubScopes) of CondOp should not be null.");
Z
zchen0211 已提交
56 57 58 59 60 61 62 63 64 65 66 67 68
  auto& sub_scopes = *sub_scopes_var->GetMutable<std::vector<Scope*>>();

  for (int i = 0; i < 2; ++i) {
    // Create two sub scopes for true and false branches
    // sub_scopes[0] for the true branch and sub_scopes[1] for the false
    // branch
    CreateScope(scope);

    // Create two tensors for true and false indices
    // index_tensors[0] for the true branch and index_tensors[1] for the false
    // branch
    CreateIndexTensor(scope);

69 70
    PADDLE_ENFORCE(!Inputs("Xs").empty(),
                   "Inputs(Xs) of CondOp can't be empty.");
Z
zchen0211 已提交
71 72 73
    for (auto& input : Inputs("Xs")) {
      // Create a new tensor in sub-scope for input-type tensor
      Variable* v = sub_scopes[i]->NewVar(input);
Z
zchen0211 已提交
74 75
      LoDTensor* sub_input = v->GetMutable<LoDTensor>();
      sub_input->Resize(scope.FindVar(input)->GetMutable<LoDTensor>()->dims());
Z
zchen0211 已提交
76 77 78 79 80 81 82 83 84
    }

    for (auto& output : (*sub_net_op_[i]).Outputs()) {
      for (auto& var_name : output.second) {
        sub_scopes[i]->NewVar(var_name);
      }
    }

    // each net calls InferShape
Y
Yu Yang 已提交
85
    //    sub_net_op_[i]->InferShape(*sub_scopes[i]);
Z
zchen0211 已提交
86 87 88
  }

  for (auto& output : Outputs("Outs")) {
Z
zchen0211 已提交
89 90 91 92 93 94
    LoDTensor* tensor_t_out =
        sub_scopes[0]->FindVar(output)->GetMutable<LoDTensor>();
    PADDLE_ENFORCE_NOT_NULL(tensor_t_out, "True output should not be NULL");
    LoDTensor* tensor_f_out =
        sub_scopes[1]->FindVar(output)->GetMutable<LoDTensor>();
    PADDLE_ENFORCE_NOT_NULL(tensor_f_out, "False output should not be NULL");
Z
zchen0211 已提交
95 96 97

    auto* tensor_out_var = scope.FindVar(output);
    PADDLE_ENFORCE_NOT_NULL(tensor_out_var, "Output not found");
Z
zchen0211 已提交
98 99 100 101
    LoDTensor* tensor_out = tensor_out_var->GetMutable<LoDTensor>();
    PADDLE_ENFORCE_NOT_NULL(tensor_t_out,
                            "True output tensor should not be NULL");

Z
zchen0211 已提交
102 103 104 105
    // check output size should be same
    PADDLE_ENFORCE_EQ(tensor_t_out->dims(), tensor_f_out->dims(),
                      "Outputs not of the same shape");
    tensor_out->Resize(tensor_t_out->dims());
Z
zchen0211 已提交
106 107 108
    // tensor_out->mutable_data<float>(tensor_out->dims(),
    // platform::CPUPlace());
    tensor_out->mutable_data<float>(platform::CPUPlace());
Z
zchen0211 已提交
109 110 111 112 113
  }
}

void CondOp::Run(const Scope& scope,
                 const platform::DeviceContext& dev_ctx) const {
Z
zchen0211 已提交
114
  auto* sub_scopes_var = scope.FindVar("SubScopes");
115 116
  PADDLE_ENFORCE_NOT_NULL(sub_scopes_var,
                          "Output(SubScopes) of CondOp should not be null.");
Z
zchen0211 已提交
117 118
  auto sub_scopes = sub_scopes_var->Get<std::vector<Scope*>>();
  auto* index_tensors_var = scope.FindVar("IndexTensors");
119 120
  PADDLE_ENFORCE_NOT_NULL(index_tensors_var,
                          "Output(IndexTensors) of CondOp should not be null.");
Z
zchen0211 已提交
121
  auto index_tensors = index_tensors_var->Get<std::vector<LoDTensor>>();
Z
zchen0211 已提交
122 123 124

  std::string cond_name = Input("Cond");
  Variable* cond_var = scope.FindVar(cond_name);
125 126
  PADDLE_ENFORCE_NOT_NULL(cond_var,
                          "Input(Cond) of CondOp should not be null.");
Z
zchen0211 已提交
127
  const LoDTensor* cond = cond_var->GetMutable<LoDTensor>();
Z
zchen0211 已提交
128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159

  // Step 1: get the true/false index at runtime
  // index_[0]: vector<int>, contains all index for cond[i] == true
  // index_[1]: vector<int>, contains all index for cond[i] == false
  for (int i = 0; i < 2; ++i) index_[i].clear();

  const int* cond_data = cond->data<int>();
  for (int i = 0; i < cond->dims()[0]; ++i) {
    if (cond_data[i])
      index_[0].push_back(i);
    else
      index_[1].push_back(i);
  }

  // put index_[0] and index_[1] into two tensors:
  // index_tensor_[0] and index_tensor_[1]
  DDim dim = paddle::framework::make_ddim({0});
  for (int i = 0; i < 2; ++i) {
    dim[0] = index_[i].size();
    int* tmp_ptr =
        index_tensors[i].mutable_data<int>(dim, platform::CPUPlace());
    index_tensors[i].Resize(dim);
    memcpy(tmp_ptr, index_[i].data(), dim[0] * sizeof(int));
  }

  // Step 2: collect data by calling gather
  for (int i = 0; i < 2; ++i) {
    // i= 0/i for True and False branches respectively
    for (auto& input : Inputs("Xs")) {
      // find Tensor
      Variable* v = scope.FindVar(input);
      PADDLE_ENFORCE_NOT_NULL(v);
Z
zchen0211 已提交
160
      LoDTensor* tensor_parent = v->GetMutable<LoDTensor>();
Z
zchen0211 已提交
161 162 163

      v = sub_scopes[i]->FindVar(input);
      PADDLE_ENFORCE_NOT_NULL(v);
Z
zchen0211 已提交
164
      LoDTensor* tensor_child = v->GetMutable<LoDTensor>();
Z
zchen0211 已提交
165 166 167 168 169 170 171 172 173 174 175 176 177

      // Resize child
      DDim dim = tensor_child->dims();
      dim[0] = index_[i].size();
      tensor_child->Resize(dim);
      tensor_child->mutable_data<float>(dim, platform::CPUPlace());

      Gather<float>(dev_ctx.GetPlace(), tensor_parent, &index_tensors[i],
                    tensor_child);
    }
  }

  // Step 3: run
Z
zchen0211 已提交
178 179 180
  for (int i = 0; i < 2; ++i) {
    sub_net_op_[i]->Run(*sub_scopes[i], dev_ctx);
  }
Z
zchen0211 已提交
181 182

  // Step 4: merge output results
183 184
  PADDLE_ENFORCE(!Outputs("Outs").empty(),
                 "Outputs(Outs) of CondOp can't be empty.");
Z
zchen0211 已提交
185 186 187 188 189 190
  for (int i = 0; i < 2; ++i) {
    // i= 0/i for True and False branches respectively
    for (auto& output : Outputs("Outs")) {
      // find Tensor
      Variable* v = scope.FindVar(output);
      PADDLE_ENFORCE_NOT_NULL(v);
Z
zchen0211 已提交
191
      LoDTensor* tensor_parent = v->GetMutable<LoDTensor>();
Z
zchen0211 已提交
192 193 194

      v = sub_scopes[i]->FindVar(output);
      PADDLE_ENFORCE_NOT_NULL(v);
Z
zchen0211 已提交
195
      LoDTensor* tensor_child = v->GetMutable<LoDTensor>();
Z
zchen0211 已提交
196 197 198 199 200 201 202 203

      ScatterUpdate<float>(dev_ctx.GetPlace(), tensor_child, &index_tensors[i],
                           tensor_parent);
    }
  }
}

class CondOpProtoAndCheckerMaker : public framework::OpProtoAndCheckerMaker {
Z
cond op  
zchen0211 已提交
204
 public:
Z
zchen0211 已提交
205 206
  CondOpProtoAndCheckerMaker(framework::OpProto* proto,
                             framework::OpAttrChecker* op_checker)
Z
cond op  
zchen0211 已提交
207 208 209 210 211 212 213 214 215 216
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput("Cond", "The condition, which is a bool vector");
    AddInput("Xs", "Inputs of Subnets").AsDuplicable();
    AddOutput("Outs", "Outputs of Cond_Op after merge").AsDuplicable();

    AddOutput("SubScopes", "sub scopes for true and false branches");
    AddOutput("IndexTensors", "Index Tensors contains indices for true/false");

    AddComment(R"DOC(
Sample dependent Cond Operator:
Z
zchen0211 已提交
217
Given Cond[i] as a 1/0 vector to indicate true/false
Q
Qiao Longfei 已提交
218
The equation is:
Z
zchen0211 已提交
219
Out[i] = subnet_t[i], if Cond[i] == true
Z
cond op  
zchen0211 已提交
220 221 222 223 224 225 226 227
Out[i] = subnet_t[i], if Cond[i] == false
)DOC");
  }
};

}  // namespace operators
}  // namespace paddle

Z
zchen0211 已提交
228
REGISTER_OP_WITHOUT_GRADIENT(cond, paddle::operators::CondOp,
Z
cond op  
zchen0211 已提交
229
                             paddle::operators::CondOpProtoAndCheckerMaker);