pybind.cc 20.5 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6

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

7
http://www.apache.org/licenses/LICENSE-2.0
8 9 10 11 12 13 14

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. */

Y
Yi Wang 已提交
15
#include "paddle/fluid/pybind/protobuf.h"
Q
qijun 已提交
16

Q
QI JUN 已提交
17
#include <mutex>  // for call_once
18
#include <unordered_map>
Y
Yi Wang 已提交
19
#include "paddle/fluid/framework/backward.h"
20
#include "paddle/fluid/framework/channel.h"
Y
Yi Wang 已提交
21 22 23 24 25 26 27
#include "paddle/fluid/framework/executor.h"
#include "paddle/fluid/framework/feed_fetch_method.h"
#include "paddle/fluid/framework/framework.pb.h"
#include "paddle/fluid/framework/init.h"
#include "paddle/fluid/framework/lod_rank_table.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/lod_tensor_array.h"
Y
Yu Yang 已提交
28
#include "paddle/fluid/framework/parallel_executor.h"
Y
Yi Wang 已提交
29
#include "paddle/fluid/framework/prune.h"
Y
Refine  
Yu Yang 已提交
30
#include "paddle/fluid/framework/reader.h"
Y
Yi Wang 已提交
31 32 33 34
#include "paddle/fluid/framework/selected_rows.h"
#include "paddle/fluid/operators/cond_op.h"
#include "paddle/fluid/operators/net_op.h"
#include "paddle/fluid/platform/enforce.h"
35
#include "paddle/fluid/platform/gpu_info.h"
Y
Yi Wang 已提交
36 37 38 39 40
#include "paddle/fluid/platform/place.h"
#include "paddle/fluid/platform/profiler.h"
#include "paddle/fluid/pybind/const_value.h"
#include "paddle/fluid/pybind/exception.h"
#include "paddle/fluid/pybind/pybind.h"
Y
Yu Yang 已提交
41
#include "paddle/fluid/pybind/recordio.h"
Y
Yi Wang 已提交
42
#include "paddle/fluid/pybind/tensor_py.h"
Y
Yu Yang 已提交
43

44
#include "paddle/fluid/string/to_string.h"
45

D
Dong Zhihong 已提交
46
#ifdef PADDLE_WITH_CUDA
Y
Yi Wang 已提交
47 48 49
#include "paddle/fluid/operators/nccl/nccl_gpu_common.h"
#include "paddle/fluid/platform/cuda_profiler.h"
#include "paddle/fluid/platform/gpu_info.h"
D
Dong Zhihong 已提交
50 51
#endif

Q
Qiao Longfei 已提交
52 53 54
// disable auto conversion to list in Python
PYBIND11_MAKE_OPAQUE(paddle::framework::LoDTensorArray);

55
namespace paddle {
56
namespace pybind {
57
bool IsCompiledWithCUDA() {
58
#ifndef PADDLE_WITH_CUDA
Q
qijun 已提交
59 60 61 62 63 64
  return false;
#else
  return true;
#endif
}

65 66
PYBIND11_PLUGIN(core) {
  py::module m("core", "C++ core of PaddlePaddle");
67

68 69 70 71
  // using framework in this function. Since it is inside a function, it will
  // not cause namespace pollution.
  using namespace paddle::framework;  // NOLINT

Y
Yu Yang 已提交
72 73
  BindException(m);

74 75 76
  py::class_<Tensor>(m, "Tensor", py::buffer_protocol())
      .def_buffer(
          [](Tensor &self) -> py::buffer_info { return CastToPyBuffer(self); })
Y
Yu Yang 已提交
77
      .def("get_dims",
78
           [](const Tensor &self) { return vectorize(self.dims()); })
Y
Yu Yang 已提交
79
      .def("set_dims",
Q
qijun 已提交
80
           [](Tensor &self, const std::vector<int64_t> &dim) {
Y
Yu Yang 已提交
81
             self.Resize(make_ddim(dim));
Y
Yu Yang 已提交
82
           })
D
dzhwinter 已提交
83 84 85 86
      .def("set_layout",
           [](Tensor &self, const std::string &layout) {
             self.set_layout(StringToDataLayout(layout));
           })
Y
Yu Yang 已提交
87
      .def("alloc_float",
D
dzhwinter 已提交
88
           [](Tensor &self, paddle::platform::CUDAPlace &place) {
Q
qijun 已提交
89
             self.mutable_data<float>(place);
Y
Yu Yang 已提交
90
           })
Q
qijun 已提交
91
      .def("alloc_float",
Y
Yu Yang 已提交
92
           [](Tensor &self, paddle::platform::CPUPlace &place) {
Q
qijun 已提交
93
             self.mutable_data<float>(place);
Y
Yu Yang 已提交
94 95
           })
      .def("alloc_int",
Y
Yu Yang 已提交
96
           [](Tensor &self, paddle::platform::CPUPlace &place) {
Q
qijun 已提交
97
             self.mutable_data<int>(place);
Y
Yu Yang 已提交
98
           })
Q
qijun 已提交
99
      .def("alloc_int",
D
dzhwinter 已提交
100
           [](Tensor &self, paddle::platform::CUDAPlace &place) {
Q
qijun 已提交
101
             self.mutable_data<int>(place);
Q
qijun 已提交
102
           })
Y
Yu Yang 已提交
103 104
      .def("set", PyCPUTensorSetFromArray<float>)
      .def("set", PyCPUTensorSetFromArray<int>)
105
      .def("set", PyCPUTensorSetFromArray<double>)
106
      .def("set", PyCPUTensorSetFromArray<int64_t>)
Y
Yu Yang 已提交
107
      .def("set", PyCPUTensorSetFromArray<bool>)
108
      .def("set", PyCPUTensorSetFromArray<uint16_t>)
109
#ifdef PADDLE_WITH_CUDA
Y
Yu Yang 已提交
110 111
      .def("set", PyCUDATensorSetFromArray<float>)
      .def("set", PyCUDATensorSetFromArray<int>)
112
      .def("set", PyCUDATensorSetFromArray<double>)
113
      .def("set", PyCUDATensorSetFromArray<int64_t>)
Y
Yu Yang 已提交
114
      .def("set", PyCUDATensorSetFromArray<bool>)
115
      .def("set", PyCUDATensorSetFromArray<uint16_t>)
Q
qijun 已提交
116
#endif
117
      .def("shape", [](Tensor &self) { return vectorize(self.dims()); })
118 119 120 121 122
      .def("set_float_element", TensorSetElement<float>)
      .def("get_float_element", TensorGetElement<float>)
      .def("set_double_element", TensorSetElement<double>)
      .def("get_double_element", TensorGetElement<double>)
      .def("dtype", [](Tensor &self) { return ToDataType(self.type()); });
Y
Yu Yang 已提交
123

124
  py::class_<LoDTensor, Tensor>(m, "LoDTensor")
125 126
      .def_buffer(
          [](Tensor &self) -> py::buffer_info { return CastToPyBuffer(self); })
127 128 129
      .def(
          "__init__",
          [](LoDTensor &instance, const std::vector<std::vector<size_t>> &lod) {
D
dzhwinter 已提交
130 131 132 133
            LoD new_lod;
            new_lod.reserve(lod.size());
            std::copy(lod.begin(), lod.end(), std::back_inserter(new_lod));
            new (&instance) LoDTensor(new_lod);
134
          })
Y
Yu Yang 已提交
135
      .def("__init__", [](LoDTensor &instance) { new (&instance) LoDTensor(); })
D
dangqingqing 已提交
136
      .def("set_lod",
137
           [](LoDTensor &self, const std::vector<std::vector<size_t>> &lod) {
Y
Yu Yang 已提交
138
             LoD new_lod;
139 140 141
             new_lod.reserve(lod.size());
             std::copy(lod.begin(), lod.end(), std::back_inserter(new_lod));
             self.set_lod(new_lod);
D
dangqingqing 已提交
142
           })
143
      .def("lod", [](LoDTensor &self) -> std::vector<std::vector<size_t>> {
D
dzhwinter 已提交
144 145 146 147 148
        auto lod = self.lod();
        std::vector<std::vector<size_t>> new_lod;
        new_lod.reserve(lod.size());
        std::copy(lod.begin(), lod.end(), std::back_inserter(new_lod));
        return new_lod;
D
dangqingqing 已提交
149 150
      });

Q
qijun 已提交
151 152 153 154 155 156 157 158 159 160 161 162 163
  py::class_<SelectedRows>(m, "SelectedRows")
      .def("__init__",
           [](SelectedRows &instance) { new (&instance) SelectedRows(); })
      .def("__init__",
           [](SelectedRows &instance, const std::vector<int64_t> rows,
              const int64_t &height) {
             new (&instance) SelectedRows(rows, height);
           })
      .def("get_tensor",
           [](SelectedRows &self) { return self.mutable_value(); },
           py::return_value_policy::reference)
      .def("set_height", &SelectedRows::set_height)
      .def("height", &SelectedRows::height)
Q
qijun 已提交
164 165 166 167 168 169 170 171 172
      .def("set_rows",
           [](SelectedRows &self, std::vector<int64_t> rows) {
#ifndef PADDLE_WITH_CUDA
             self.set_rows(rows);
#else
        Vector<int64_t> new_rows(rows);
        self.set_rows(new_rows);
#endif
           })
173 174 175 176 177 178 179 180 181 182 183
      .def("rows", [](SelectedRows &self) {
#ifndef PADDLE_WITH_CUDA
        return self.rows();
#else
         auto rows = self.rows();
         std::vector<int64_t> new_rows;
         new_rows.reserve(rows.size());
         std::copy(rows.begin(), rows.end(), std::back_inserter(new_rows));
         return new_rows;
#endif
      });
Q
qijun 已提交
184

185
  py::class_<Variable>(m, "Variable", R"DOC(Variable Class.
186 187 188

All parameter, weight, gradient are variables in Paddle.
)DOC")
189
      .def("is_int", [](const Variable &var) { return var.IsType<int>(); })
190
      .def("set_int",
191 192
           [](Variable &var, int val) -> void { *var.GetMutable<int>() = val; })
      .def("get_int", [](const Variable &var) -> int { return var.Get<int>(); })
193 194 195 196 197 198 199
      .def("is_float", [](const Variable &var) { return var.IsType<float>(); })
      .def("set_float",
           [](Variable &var, float val) -> void {
             *var.GetMutable<float>() = val;
           })
      .def("get_float",
           [](const Variable &var) -> float { return var.Get<float>(); })
Y
Yu Yang 已提交
200
      .def("get_tensor",
201 202
           [](Variable &self) -> LoDTensor * {
             return self.GetMutable<LoDTensor>();
D
dangqingqing 已提交
203 204
           },
           py::return_value_policy::reference)
Y
Yu Yang 已提交
205 206 207
      .def("get_lod_rank_table",
           [](Variable &self) { return self.GetMutable<LoDRankTable>(); },
           py::return_value_policy::reference)
Q
qijun 已提交
208 209 210 211 212
      .def("get_selected_rows",
           [](Variable &self) -> SelectedRows * {
             return self.GetMutable<SelectedRows>();
           },
           py::return_value_policy::reference)
Y
Yu Yang 已提交
213 214 215
      .def("get_lod_tensor_array",
           [](Variable &self) { return self.GetMutable<LoDTensorArray>(); },
           py::return_value_policy::reference)
D
Dong Zhihong 已提交
216 217 218 219 220 221 222
#ifdef PADDLE_WITH_CUDA
      .def("get_communicator",
           [](Variable &self) -> platform::Communicator * {
             return self.GetMutable<platform::Communicator>();
           },
           py::return_value_policy::reference)
#endif
Y
Yan Chunwei 已提交
223
      .def("get_net",
D
dongzhihong 已提交
224 225
           [](Variable &self) -> operators::NetOp * {
             return self.GetMutable<operators::NetOp>();
Y
Yan Chunwei 已提交
226
           },
Y
Refine  
Yu Yang 已提交
227 228 229 230 231 232
           py::return_value_policy::reference)
      .def("get_reader",
           [](Variable &self) -> framework::ReaderHolder * {
             PADDLE_ENFORCE(self.IsType<framework::ReaderHolder>());
             return self.GetMutable<framework::ReaderHolder>();
           },
Y
Yu Yang 已提交
233
           py::return_value_policy::reference);
234

Y
Refine  
Yu Yang 已提交
235 236 237 238
  py::class_<framework::ReaderHolder>(m, "Reader", "")
      .def("has_next", &framework::ReaderHolder::HasNext)
      .def("reset", &framework::ReaderHolder::ReInit);

239
  py::class_<Scope>(m, "Scope", "")
D
dongzhihong 已提交
240
      .def("var",
241
           [](Scope &self, const std::string &name) -> Variable * {
D
dongzhihong 已提交
242
             return self.Var(name);
Y
Yu Yang 已提交
243
           },
244
           py::return_value_policy::reference)
245
      .def("find_var", &Scope::FindVar, py::return_value_policy::reference)
Y
Yu Yang 已提交
246
      .def(py::init<>())
247
      .def("new_scope", [](Scope &self) -> Scope * { return &self.NewScope(); },
248
           py::return_value_policy::reference)
Y
Yu Yang 已提交
249
      .def("drop_kids", &Scope::DropKids);
250

Y
Yu Yang 已提交
251 252
  //! @note: Be careful! PyBind will return std::string as an unicode, not
  //! Python str. If you want a str object, you should cast them in Python.
Y
Yu Yang 已提交
253 254
  m.def("get_all_op_protos", []() -> std::vector<py::bytes> {
    std::vector<py::bytes> ret_values;
255 256 257 258 259 260 261 262 263 264
    for (auto &iter : OpInfoMap::Instance().map()) {
      auto &info = iter.second;
      if (info.HasOpProtoAndChecker()) {
        std::string str;
        PADDLE_ENFORCE(
            info.Proto().SerializeToString(&str),
            "Serialize OpProto Error. This could be a bug of Paddle.");
        ret_values.emplace_back(str);
      }
    }
Y
Yu Yang 已提交
265 266
    return ret_values;
  });
267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282
  m.def(
      "get_grad_op_desc", [](const OpDesc &op_desc,
                             const std::unordered_set<std::string> &no_grad_set,
                             const std::vector<BlockDesc *> &grad_sub_block) {
        std::unordered_map<std::string, std::string> grad_to_var;
        std::vector<std::unique_ptr<OpDesc>> grad_op_descs =
            framework::OpInfoMap::Instance()
                .Get(op_desc.Type())
                .GradOpMaker()(op_desc, no_grad_set, &grad_to_var,
                               grad_sub_block);
        std::vector<OpDesc *> grad_op_desc_ptrs(grad_op_descs.size());
        std::transform(grad_op_descs.begin(), grad_op_descs.end(),
                       grad_op_desc_ptrs.begin(),
                       [](std::unique_ptr<OpDesc> &p) { return p.release(); });
        return std::make_pair(grad_op_desc_ptrs, grad_to_var);
      });
Y
Yu Yang 已提交
283
  m.def("prune", [](const ProgramDesc &origin,
284
                    const std::vector<std::array<size_t, 2>> &targets) {
Y
Yu Yang 已提交
285
    ProgramDesc prog_with_targets(origin);
286
    for (const auto &t : targets) {
287
      prog_with_targets.MutableBlock(t[0])->Op(t[1])->MarkAsTarget();
288
    }
289
    proto::ProgramDesc pruned_desc;
290
    Prune(*prog_with_targets.Proto(), &pruned_desc);
Y
Yu Yang 已提交
291
    return new ProgramDesc(pruned_desc);
292
  });
Y
Yu Yang 已提交
293
  m.def("inference_optimize", [](ProgramDesc &origin) {
294
    proto::ProgramDesc pruned_desc;
295
    InferenceOptimize(*(origin.Proto()), &pruned_desc);
Y
Yu Yang 已提交
296
    return new ProgramDesc(pruned_desc);
297
  });
F
fengjiayi 已提交
298 299
  m.def("empty_var_name", []() { return framework::kEmptyVarName; });
  m.def("grad_var_suffix", []() { return framework::kGradVarSuffix; });
300 301 302
  m.def_submodule(
       "var_names",
       "The module will return special predefined variable name in Paddle")
Y
Yi Wang 已提交
303 304
      .def("empty", []() { return kEmptyVarName; })
      .def("temp", []() { return kTempVarName; });
Q
qijun 已提交
305
  // clang-format off
Y
Yu Yang 已提交
306
  py::class_<paddle::platform::DeviceContext>(m, "DeviceContext")
Q
qijun 已提交
307 308
      .def_static("create",
                  [](paddle::platform::CPUPlace& place)
Q
qijun 已提交
309
                      -> paddle::platform::DeviceContext* {
Q
qijun 已提交
310 311 312
                    return new paddle::platform::CPUDeviceContext();
                  })
      .def_static("create",
D
dzhwinter 已提交
313
                  [](paddle::platform::CUDAPlace& place)
Q
qijun 已提交
314
                      -> paddle::platform::DeviceContext* {
315
#ifndef PADDLE_WITH_CUDA
D
dzhwinter 已提交
316
                    PADDLE_THROW("CUDAPlace is not supported in CPU device.");
Q
qijun 已提交
317
#else
Q
qijun 已提交
318
                    return new paddle::platform::CUDADeviceContext(place);
Q
qijun 已提交
319
#endif
Q
qijun 已提交
320
                  });
D
Dong Zhihong 已提交
321 322 323 324
// clang-format on
#ifdef PADDLE_WITH_CUDA
  py::class_<platform::Communicator>(m, "Communicator").def(py::init<>());
#endif
D
dzhwinter 已提交
325
  py::class_<platform::CUDAPlace>(m, "CUDAPlace")
326
      .def(py::init<int>())
D
dzhwinter 已提交
327
      .def("__str__", string::to_string<const platform::CUDAPlace &>);
Q
qijun 已提交
328

329 330 331
  py::class_<paddle::platform::CPUPlace>(m, "CPUPlace")
      .def(py::init<>())
      .def("__str__", string::to_string<const platform::CPUPlace &>);
Y
Yu Yang 已提交
332

Y
Yu Yang 已提交
333 334 335 336 337 338 339
  py::class_<platform::Place>(m, "Place")
      .def(py::init<>())
      .def("set_place",
           [](platform::Place &self, const platform::CPUPlace &cpu_place) {
             self = cpu_place;
           })
      .def("set_place",
D
dzhwinter 已提交
340
           [](platform::Place &self, const platform::CUDAPlace &gpu_place) {
Y
Yu Yang 已提交
341 342 343
             self = gpu_place;
           });

Y
Yu Yang 已提交
344 345 346
  py::class_<OperatorBase>(m, "Operator")
      .def_static("create",
                  [](py::bytes protobin) {
347
                    proto::OpDesc desc;
Y
Yu Yang 已提交
348 349 350 351 352
                    PADDLE_ENFORCE(desc.ParsePartialFromString(protobin),
                                   "Cannot parse user input to OpDesc");
                    PADDLE_ENFORCE(desc.IsInitialized(),
                                   "User OpDesc is not initialized, reason %s",
                                   desc.InitializationErrorString());
353
                    return OpRegistry::CreateOp(desc);
Y
Yu Yang 已提交
354 355 356 357 358 359
                  })
      .def("backward",
           [](const OperatorBase &forwardOp,
              const std::unordered_set<std::string> &no_grad_vars) {
             return Backward(forwardOp, no_grad_vars).release();
           })
360
      .def("run",
361
           [](OperatorBase &self, const Scope &scope,
D
dzhwinter 已提交
362 363 364
              const platform::CPUPlace &place) { self.Run(scope, place); })
      .def("run",
           [](OperatorBase &self, const Scope &scope,
D
dzhwinter 已提交
365
              const platform::CUDAPlace &place) { self.Run(scope, place); })
Y
Yu Yang 已提交
366 367 368 369 370 371 372
      .def("type",
           [](const OperatorBase &op) -> std::string { return op.Type(); })
      .def("outputs",
           [](const OperatorBase &op)
               -> std::map<std::string, std::vector<std::string>> {
                 return op.Outputs();
               })
Q
qijun 已提交
373 374
      .def("output_vars",
           [](const OperatorBase &op) { return op.OutputVars(true); })
Y
Yu Yang 已提交
375
      .def("inputs", [](const OperatorBase &op) { return op.Inputs(); })
Q
qijun 已提交
376
      .def("input_vars", [](const OperatorBase &op) { return op.InputVars(); })
Y
Yu Yang 已提交
377 378 379 380
      .def("__str__", &OperatorBase::DebugString)
      .def("no_intermediate_outputs",
           [](const OperatorBase &op) { return op.OutputVars(false); })
      .def("support_gpu", &OperatorBase::SupportGPU);
Y
Yu Yang 已提交
381

Y
Yu Yang 已提交
382 383 384 385 386 387 388
  py::class_<operators::NetOp, OperatorBase>(m, "Net")
      .def_static("create",
                  []() -> operators::NetOp * {
                    auto *retv = new operators::NetOp;
                    retv->SetType("plain_net");
                    return retv;
                  })
389 390
      .def("append_op", [](operators::NetOp &self,
                           const OperatorBase &op) { self.AppendOp(op); })
D
dongzhihong 已提交
391 392 393 394
      .def("complete_add_op", &operators::NetOp::CompleteAddOp)
      .def("complete_add_op", [](std::shared_ptr<operators::NetOp> &self) {
        self->CompleteAddOp();
      });
Y
Yan Chunwei 已提交
395

Z
cond op  
zchen0211 已提交
396 397 398 399
  // cond_op
  py::class_<operators::CondOp, OperatorBase>(m, "CondOp")
      .def_static("create",
                  [](py::bytes protobin) -> operators::CondOp * {
400
                    proto::OpDesc desc;
Z
cond op  
zchen0211 已提交
401 402 403 404 405
                    PADDLE_ENFORCE(desc.ParsePartialFromString(protobin),
                                   "Cannot parse user input to OpDesc");
                    PADDLE_ENFORCE(desc.IsInitialized(),
                                   "User OpDesc is not initialized, reason %s",
                                   desc.InitializationErrorString());
406
                    auto cond_op = OpRegistry::CreateOp(desc);
Z
cond op  
zchen0211 已提交
407 408 409 410 411 412 413 414 415 416 417
                    return static_cast<operators::CondOp *>(cond_op.release());
                  })
      .def("set_truenet",
           [](operators::CondOp &self, const operators::NetOp &net) -> void {
             self.set_truenet(net.Clone());
           })
      .def("set_falsenet",
           [](operators::CondOp &self, const operators::NetOp &net) -> void {
             self.set_falsenet(net.Clone());
           });

F
fengjiayi 已提交
418
  py::class_<framework::Executor>(m, "Executor")
D
dzhwinter 已提交
419
      .def(py::init<const platform::Place &>())
420 421 422
      .def("run",
           (void (Executor::*)(const ProgramDesc &, Scope *, int, bool, bool)) &
               Executor::Run);
F
fengjiayi 已提交
423

D
dzhwinter 已提交
424
  m.def("init_gflags", framework::InitGflags);
Y
Yang Yu 已提交
425
  m.def("init_glog", framework::InitGLOG);
D
dzhwinter 已提交
426
  m.def("init_devices", &framework::InitDevices);
427

428
  m.def("is_compiled_with_cuda", IsCompiledWithCUDA);
429 430 431 432 433 434
#ifdef PADDLE_WITH_CUDA
  m.def("is_float16_supported", [](const platform::CUDAPlace &place) -> bool {
    // Only GPUs with Compute Capability >= 53 support float16
    return platform::GetCUDAComputeCapability(place.device) >= 53;
  });
#endif
435

436
  m.def("set_feed_variable", framework::SetFeedVariable);
Q
qijun 已提交
437
  m.def("get_fetch_variable", framework::GetFetchVariable);
Q
qijun 已提交
438

F
fengjiayi 已提交
439 440 441 442
  BindProgramDesc(m);
  BindBlockDesc(m);
  BindVarDsec(m);
  BindOpDesc(m);
Q
qiaolongfei 已提交
443
  BindConstValue(m);
Y
Yu Yang 已提交
444

Y
Yu Yang 已提交
445 446 447 448 449 450 451 452 453
  py::class_<framework::LoDRankTable>(m, "LodRankTable")
      .def("items", [](framework::LoDRankTable &table) {
        std::vector<std::pair<size_t, size_t>> res;
        for (auto &item : table.items()) {
          res.push_back({item.index, item.length});
        }
        return res;
      });

Y
Yu Yang 已提交
454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470
  py::class_<LoDTensorArray>(m, "LoDTensorArray")
      .def("__getitem__",
           [](LoDTensorArray &self, size_t i) { return &self.at(i); },
           py::return_value_policy::reference)
      .def("__len__", [](LoDTensorArray &self) { return self.size(); })
      .def("__setitem__",
           [](LoDTensorArray &self, size_t i, const LoDTensor &t) {
             PADDLE_ENFORCE_LT(i, self.size());
             self[i].ShareDataWith(t);
             self[i].set_lod(t.lod());
           })
      .def("append", [](LoDTensorArray &self, const LoDTensor &t) {
        self.emplace_back();
        self.back().ShareDataWith(t);
        self.back().set_lod(t.lod());
      });

Y
Yu Yang 已提交
471
  m.def("op_support_gpu", OpSupportGPU);
D
Dong Zhihong 已提交
472
#ifdef PADDLE_WITH_CUDA
D
Dong Zhihong 已提交
473
  m.def("get_cuda_device_count", platform::GetCUDADeviceCount);
D
dangqingqing 已提交
474 475 476 477

  m.def("nvprof_init", platform::CudaProfilerInit);
  m.def("nvprof_start", platform::CudaProfilerStart);
  m.def("nvprof_stop", platform::CudaProfilerStop);
D
Dong Zhihong 已提交
478
#endif
Y
Yu Yang 已提交
479

480 481 482 483
  py::enum_<platform::ProfilerState>(m, "ProfilerState", py::arithmetic())
      .value("kDisabled", platform::ProfilerState::kDisabled)
      .value("kCPU", platform::ProfilerState::kCPU)
      .value("kCUDA", platform::ProfilerState::kCUDA)
484
      .value("kAll", platform::ProfilerState::kAll)
485 486 487 488 489 490 491 492 493 494 495 496 497 498
      .export_values();

  py::enum_<platform::EventSortingKey>(m, "EventSortingKey", py::arithmetic())
      .value("kDefault", platform::EventSortingKey::kDefault)
      .value("kCalls", platform::EventSortingKey::kCalls)
      .value("kTotal", platform::EventSortingKey::kTotal)
      .value("kMin", platform::EventSortingKey::kMin)
      .value("kMax", platform::EventSortingKey::kMax)
      .value("kAve", platform::EventSortingKey::kAve)
      .export_values();

  m.def("enable_profiler", platform::EnableProfiler);
  m.def("disable_profiler", platform::DisableProfiler);
  m.def("reset_profiler", platform::ResetProfiler);
Y
Yu Yang 已提交
499

Y
Yu Yang 已提交
500 501 502 503 504 505 506 507 508 509 510 511 512
  py::class_<ParallelExecutor>(m, "ParallelExecutor")
      .def(
          "__init__",
          [](ParallelExecutor &self, const std::vector<platform::Place> &places,
             const std::unordered_set<std::string> &params,
             const ProgramDesc &startup_program,
             const ProgramDesc &main_program, const std::string &loss_var_name,
             Scope *scope) {
            new (&self) ParallelExecutor(places, params, startup_program,
                                         main_program, loss_var_name, scope);
          })
      .def("run", [](ParallelExecutor &self) { self.Run({}); });

Y
Yu Yang 已提交
513
  BindRecordIOWriter(m);
514
  return m.ptr();
L
Luo Tao 已提交
515
}
516
}  // namespace pybind
517
}  // namespace paddle