tensor_py.h 3.9 KB
Newer Older
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. */

#pragma once
Q
qijun 已提交
16
#include <string>
Q
qijun 已提交
17 18 19 20
#include "paddle/framework/tensor.h"
#include "paddle/memory/memcpy.h"
#include "pybind11/numpy.h"
#include "pybind11/pybind11.h"
21 22 23 24 25

namespace py = pybind11;

namespace paddle {

26
namespace pybind {
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44

namespace details {

template <bool less, size_t I, typename... ARGS>
struct CastToPyBufferImpl;

template <size_t I, typename... ARGS>
struct CastToPyBufferImpl<false, I, ARGS...> {
  py::buffer_info operator()(framework::Tensor &tensor) {
    PADDLE_THROW("This type of tensor cannot be expose to Python");
    return py::buffer_info();
  }
};

template <size_t I, typename... ARGS>
struct CastToPyBufferImpl<true, I, ARGS...> {
  using CUR_TYPE = typename std::tuple_element<I, std::tuple<ARGS...>>::type;
  py::buffer_info operator()(framework::Tensor &tensor) {
45
    if (std::type_index(typeid(CUR_TYPE)) == tensor.holder_->type()) {
46 47 48 49 50 51 52 53 54 55 56 57
      auto dim_vec = framework::vectorize(tensor.dims());
      std::vector<size_t> dims_outside;
      std::vector<size_t> strides;
      dims_outside.resize(dim_vec.size());
      strides.resize(dim_vec.size());

      size_t prod = 1;
      for (size_t i = dim_vec.size(); i != 0; --i) {
        dims_outside[i - 1] = (size_t)dim_vec[i - 1];
        strides[i - 1] = sizeof(CUR_TYPE) * prod;
        prod *= dims_outside[i - 1];
      }
Q
qijun 已提交
58
      framework::Tensor dst_tensor;
Q
qijun 已提交
59
      if (paddle::platform::is_gpu_place(tensor.holder_->place())) {
Q
qijun 已提交
60
        dst_tensor.CopyFrom<CUR_TYPE>(tensor, platform::CPUPlace());
Q
qijun 已提交
61
      } else if (paddle::platform::is_cpu_place(tensor.holder_->place())) {
Q
qijun 已提交
62 63
        dst_tensor = tensor;
      }
64
      return py::buffer_info(
Q
qijun 已提交
65
          dst_tensor.mutable_data<CUR_TYPE>(dst_tensor.holder_->place()),
66 67
          sizeof(CUR_TYPE), py::format_descriptor<CUR_TYPE>::format(),
          (size_t)framework::arity(dst_tensor.dims()), dims_outside, strides);
68 69 70 71 72 73 74 75 76 77 78 79 80
    } else {
      constexpr bool less = I + 1 < std::tuple_size<std::tuple<ARGS...>>::value;
      return CastToPyBufferImpl<less, I + 1, ARGS...>()(tensor);
    }
  }
};
}  // namespace details
inline py::buffer_info CastToPyBuffer(framework::Tensor &tensor) {
  auto buffer_info = details::CastToPyBufferImpl<true, 0, float, int>()(tensor);
  return buffer_info;
}

template <typename T>
Q
qijun 已提交
81
void PyCPUTensorSetFromArray(
82
    framework::Tensor &self,
Q
qijun 已提交
83 84
    py::array_t<T, py::array::c_style | py::array::forcecast> array,
    paddle::platform::CPUPlace &place) {
Q
qijun 已提交
85
  std::vector<int64_t> dims;
86 87 88 89 90
  dims.reserve(array.ndim());
  for (size_t i = 0; i < array.ndim(); ++i) {
    dims.push_back((int)array.shape()[i]);
  }

F
fengjiayi 已提交
91
  self.Resize(framework::make_ddim(dims));
Q
qijun 已提交
92
  auto *dst = self.mutable_data<T>(place);
93 94 95
  std::memcpy(dst, array.data(), sizeof(T) * array.size());
}

Q
qijun 已提交
96
#ifndef PADDLE_ONLY_CPU
Q
qijun 已提交
97 98 99 100 101
template <typename T>
void PyCUDATensorSetFromArray(
    framework::Tensor &self,
    py::array_t<T, py::array::c_style | py::array::forcecast> array,
    paddle::platform::GPUPlace &place) {
Q
qijun 已提交
102
  std::vector<int64_t> dims;
Q
qijun 已提交
103 104 105
  dims.reserve(array.ndim());
  for (size_t i = 0; i < array.ndim(); ++i) {
    dims.push_back((int)array.shape()[i]);
Q
qijun 已提交
106
  }
Q
qijun 已提交
107 108 109

  self.Resize(framework::make_ddim(dims));
  auto *dst = self.mutable_data<T>(place);
110 111
  paddle::platform::GpuMemcpySync(dst, array.data(), sizeof(T) * array.size(),
                                  cudaMemcpyHostToDevice);
112
}
Q
qijun 已提交
113
#endif
114 115 116

}  // namespace pybind
}  // namespace paddle