Dragon - C++ API
A Computation Graph Virtual Machine Based Deep Learning Framework
py_dragon.h
Go to the documentation of this file.
1 
13 #ifndef DRAGON_PYTHON_PY_DRAGON_H_
14 #define DRAGON_PYTHON_PY_DRAGON_H_
15 
16 #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
17 
18 #include "py_types.h"
19 
20 #include "core/common.h"
21 #include "core/context.h"
22 #include "core/operator.h"
23 #include "core/registry.h"
24 #include "core/workspace.h"
25 #include "core/context_cuda.h"
26 #include "core/graph_gradient.h"
27 #include "core/operator_gradient.h"
28 #include "utils/caffemodel.h"
29 #include "onnx/onnx_backend.h"
30 
31 #include <pybind11/stl.h>
32 #include <pybind11/pybind11.h>
33 
34 namespace dragon {
35 
36 namespace python {
37 
39  public:
40  virtual ~TensorFetcherBase() {}
41  virtual pybind11::object Fetch(const Tensor& tensor) = 0;
42 };
43 
45  public:
46  virtual ~TensorFeederBase() {}
47  virtual void Feed(
48  const DeviceOption& option,
49  PyArrayObject* array,
50  Tensor* tensor) = 0;
51 };
52 
53 DECLARE_TYPED_REGISTRY(TensorFetcherRegistry, TypeId, TensorFetcherBase);
54 
55 #define REGISTER_TENSOR_FETCHER(type, ...) \
56  REGISTER_TYPED_CLASS(TensorFetcherRegistry, type, __VA_ARGS__)
57 
59  return TensorFetcherRegistry()->Create(type);
60 }
61 
62 DECLARE_TYPED_REGISTRY(TensorFeederRegistry, TypeId, TensorFeederBase);
63 
64 #define REGISTER_TENSOR_FEEDER(type, ...) \
65  REGISTER_TYPED_CLASS(TensorFeederRegistry, type, __VA_ARGS__)
66 
68  public:
69  pybind11::object Fetch(const Tensor& tensor) override {
70  CHECK_GT(tensor.count(), 0);
71  vector<npy_intp> npy_dims;
72  for (auto dim : tensor.dims()) npy_dims.push_back(dim);
73  int npy_type = TypeMetaToNPY(tensor.meta());
74  CHECK(npy_type != -1)
75  << "\nThe data type of Tensor(" << tensor.name()
76  << ") is unknown. Have you solved it?";
77  CHECK(tensor.memory()) << "\nIllegal memory access.";
78  // Create a empty array with the same shape
79  auto* array = PyArray_SimpleNew(
80  tensor.ndim(), npy_dims.data(), npy_type);
81  // Copy the tensor data to the numpy array
82  if (tensor.memory_state() == MixedMemory::STATE_AT_CUDA) {
83  CUDAContext::MemcpyEx<CPUContext, CUDAContext>(
84  tensor.nbytes(),
85  PyArray_DATA(reinterpret_cast<PyArrayObject*>(array)),
86  tensor.raw_data<CUDAContext>(),
87  tensor.memory()->device_id()
88  );
89  } else {
90  CPUContext::Memcpy<CPUContext, CPUContext>(
91  tensor.nbytes(),
92  PyArray_DATA(reinterpret_cast<PyArrayObject*>(array)),
93  tensor.raw_data<CPUContext>()
94  );
95  }
96  return pybind11::reinterpret_steal<pybind11::object>(array);
97  }
98 };
99 
101  public:
102  pybind11::object Fetch(const Tensor& tensor) override {
103  CHECK_EQ(tensor.count(), 1);
104  return pybind11::bytes(tensor.data<string, CPUContext>()[0]);
105  }
106 };
107 
109  public:
110  void Feed(
111  const DeviceOption& option,
112  PyArrayObject* original_array,
113  Tensor* tensor) override {
114  auto* array = PyArray_GETCONTIGUOUS(original_array);
115  const auto& meta = TypeNPYToMeta(PyArray_TYPE(array));
116  if (meta.id() == 0) LOG(FATAL) << "Unsupported DType.";
117  tensor->SetMeta(meta);
118  int ndim = PyArray_NDIM(array);
119  vec64_t dims(ndim);
120  auto* npy_dims = PyArray_DIMS(array);
121  for (int i = 0; i < ndim; i++) dims[i] = npy_dims[i];
122  tensor->Reshape(dims);
123  if (option.device_type() == PROTO_CUDA) {
124 #ifdef WITH_CUDA
125  CUDAContext::MemcpyEx<CUDAContext, CPUContext>(
126  tensor->nbytes(),
127  tensor->raw_mutable_data<CUDAContext>(),
128  static_cast<void*>(PyArray_DATA(array)),
129  option.device_id()
130  );
131 #else
132  LOG(FATAL) << "CUDA was not compiled.";
133 #endif
134  } else {
135  CPUContext::Memcpy<CPUContext, CPUContext>(
136  tensor->nbytes(),
137  tensor->raw_mutable_data<CPUContext>(),
138  static_cast<void*>(PyArray_DATA(array))
139  );
140  }
141  Py_XDECREF(array);
142  }
143 };
144 
145 } // namespace python
146 
147 } // namespace dragon
148 
149 #endif // DRAGON_PYTHON_PY_DRAGON_H_
Definition: logging.h:21
virtual ~TensorFetcherBase()
Definition: py_dragon.h:40
int64_t count(int64_t start, int64_t end) const
Return the number of elements along the [start, end) axes.
Definition: tensor.h:100
intptr_t TypeId
Definition: typeid.h:22
int ndim() const
Return the number of dimensions.
Definition: tensor.h:82
Definition: py_dragon.h:67
MixedMemory::State memory_state() const
Return the state of the internal memory.
Definition: tensor.h:151
const vec64_t & dims() const
Return all the dimensions.
Definition: tensor.h:88
Memory could be modified by CUDAContext last time.
Definition: mixedmem.h:35
const TypeMeta & TypeNPYToMeta(int npy_type)
Definition: py_types.h:40
virtual pybind11::object Fetch(const Tensor &tensor)=0
MixedMemory * memory() const
Return the pointer of internal memory.
Definition: tensor.h:143
size_t nbytes() const
Return the total number bytes of this tensor.
Definition: tensor.h:94
Definition: tensor.h:21
int device_id() const
Return the device id of this context.
Definition: context_cuda.h:257
Definition: context.h:20
const TypeMeta & meta() const
Return the type meta of this tensor.
Definition: tensor.h:164
std::vector< int64_t > vec64_t
Definition: types.h:25
Definition: py_dragon.h:108
DECLARE_TYPED_REGISTRY(TensorFetcherRegistry, TypeId, TensorFetcherBase)
#define CHECK_GT(val1, val2)
Definition: logging.h:50
const T * data() const
Get the typed const data pointer.
Definition: tensor.h:282
void SetMeta(const TypeMeta &meta)
Set the type meta.
Definition: tensor.h:167
const void * raw_data() const
Get the raw const data pointer.
Definition: tensor.h:256
Tensor * Reshape(const vec64_t &dims)
Reshape to the given dimensions.
Definition: tensor.h:38
void Feed(const DeviceOption &option, PyArrayObject *original_array, Tensor *tensor) override
Definition: py_dragon.h:110
pybind11::object Fetch(const Tensor &tensor) override
Definition: py_dragon.h:69
const int TypeMetaToNPY(const TypeMeta &meta)
Definition: py_types.h:25
pybind11::object Fetch(const Tensor &tensor) override
Definition: py_dragon.h:102
Definition: py_dragon.h:44
Definition: context_cuda.h:142
Definition: py_dragon.h:100
#define LOG(severity)
Definition: logging.h:54
virtual void Feed(const DeviceOption &option, PyArrayObject *array, Tensor *tensor)=0
TensorFetcherBase * CreateFetcher(TypeId type)
Definition: py_dragon.h:58
int device_id() const
Return the device id of the memory on device.
Definition: mixedmem.h:112
Definition: py_dragon.h:38
#define CHECK(condition)
Definition: logging.h:45
virtual ~TensorFeederBase()
Definition: py_dragon.h:46
const string & name() const
Return the tensor name.
Definition: tensor.h:69
Definition: common.h:41
void * raw_mutable_data(const TypeMeta &meta)
Try to allocate the raw data memory.
Definition: tensor.h:221
#define CHECK_EQ(val1, val2)
Definition: logging.h:48