mxnet
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros
ndarray.h
Go to the documentation of this file.
1 
6 #ifndef MXNET_NDARRAY_H_
7 #define MXNET_NDARRAY_H_
8 
9 #include <dmlc/base.h>
10 #include <dmlc/logging.h>
11 #include <dmlc/io.h>
12 #include <dmlc/type_traits.h>
13 #include <dmlc/registry.h>
14 #include <vector>
15 #include <map>
16 #include <string>
17 #include <memory>
18 #include "./base.h"
19 #include "./storage.h"
20 #include "./engine.h"
21 
22 // check c++11
23 #if DMLC_USE_CXX11 == 0
24 #error "cxx11 was required for ndarray module"
25 #endif
26 
27 namespace mxnet {
31 class NDArray {
32  public:
34  NDArray() {}
43  bool delay_alloc = false, int dtype = mshadow::default_type_flag)
44  : ptr_(std::make_shared<Chunk>(shape.Size(), ctx, delay_alloc, dtype)),
45  shape_(shape), offset_(0), dtype_(dtype) {
46  }
54  NDArray(const TBlob &data, int dev_id)
55  : ptr_(std::make_shared<Chunk>(data, dev_id)), shape_(data.shape_), offset_(0),
56  dtype_(data.type_flag_) {
57  }
61  inline const TShape &shape() const {
62  return shape_;
63  }
67  inline TBlob data() const {
68  MSHADOW_TYPE_SWITCH(dtype_, DType, {
69  return TBlob(static_cast<DType*>(ptr_->shandle.dptr)
70  + offset_, shape_, ptr_->shandle.ctx.dev_mask());
71  });
72  return TBlob();
73  }
77  inline Context ctx() const {
78  return ptr_->shandle.ctx;
79  }
83  inline int dtype() const {
84  return dtype_;
85  }
87  inline bool is_none() const {
88  return ptr_.get() == nullptr;
89  }
94  inline void WaitToRead() const {
95  if (is_none()) return;
96  Engine::Get()->WaitForVar(ptr_->var);
97  }
102  inline void WaitToWrite() const {
103  if (is_none()) return;
108  Engine::Get()->PushSync([](RunContext) {}, Context{}, {}, {ptr_->var});
109  Engine::Get()->WaitForVar(ptr_->var);
110  }
112  inline Engine::VarHandle var() const {
113  return ptr_->var;
114  }
119  void Save(dmlc::Stream *strm) const;
125  bool Load(dmlc::Stream *strm);
131  NDArray &operator=(real_t scalar);
138  NDArray &operator+=(const NDArray &src);
145  NDArray &operator+=(const real_t &src);
152  NDArray &operator-=(const NDArray &src);
159  NDArray &operator-=(const real_t &src);
166  NDArray &operator*=(const NDArray &src);
173  NDArray &operator*=(const real_t &src);
180  NDArray &operator/=(const NDArray &src);
187  NDArray &operator/=(const real_t &src);
192  NDArray T() const;
198  NDArray Copy(Context ctx) const;
209  void SyncCopyFromCPU(const void *data, size_t size) const;
220  void SyncCopyToCPU(void *data, size_t size) const;
227  inline NDArray Slice(index_t begin, index_t end) const {
228  NDArray ret = *this;
229  CHECK(!is_none()) << "NDArray is not initialized";
230  CHECK_GE(shape_[0], end) << "Slice end index out of range";
231  size_t length = shape_.ProdShape(1, shape_.ndim());
232  ret.offset_ += begin * length;
233  ret.shape_[0] = end - begin;
234  return ret;
235  }
241  inline NDArray At(index_t idx) const {
242  NDArray ret = *this;
243  CHECK(!is_none()) << "NDArray is not initialized";
244  CHECK_GE(shape_[0], idx) << "index out of range";
245  size_t length = shape_.ProdShape(1, shape_.ndim());
246  ret.offset_ += idx * length;
247  ret.shape_ = TShape(shape_.data()+1, shape_.data()+shape_.ndim());
248  return ret;
249  }
255  inline NDArray Reshape(const TShape &shape) const {
256  CHECK_GE(shape_.Size(), shape.Size())
257  << "NDArray.Reshape: target shape size is different from current shape";
258  NDArray ret = *this;
259  ret.shape_ = shape;
260  return ret;
261  }
266  inline void CheckAndAlloc() const {
267  ptr_->CheckAndAlloc();
268  }
275  static void Save(dmlc::Stream* fo,
276  const std::vector<NDArray>& data,
277  const std::vector<std::string>& names);
284  static void Load(dmlc::Stream* fi,
285  std::vector<NDArray>* data,
286  std::vector<std::string>* keys);
287 
288  private:
290  struct Chunk {
292  Storage::Handle shandle;
299  bool static_data;
301  bool delay_alloc;
303  Chunk() : static_data(true), delay_alloc(false) {
304  var = Engine::Get()->NewVariable();
305  }
307  Chunk(const TBlob &data, int dev_id)
308  : static_data(true),
309  delay_alloc(false) {
310  var = Engine::Get()->NewVariable();
311  if (data.dev_mask_ == cpu::kDevMask) {
312  shandle.ctx = Context::CPU();
313  } else {
314  CHECK_EQ(data.dev_mask_, gpu::kDevMask);
315  shandle.ctx = Context::GPU(dev_id);
316  }
317  shandle.dptr = data.dptr_;
318  shandle.size = data.shape_.Size() * mshadow::mshadow_sizeof(data.type_flag_);
319  }
321  Chunk(uint64_t size, Context ctx, bool delay_alloc_, int dtype)
322  : static_data(false), delay_alloc(true) {
323  var = Engine::Get()->NewVariable();
324  shandle.size = size * mshadow::mshadow_sizeof(dtype);
325  shandle.ctx = ctx;
326  if (!delay_alloc_) this->CheckAndAlloc();
327  }
329  inline void CheckAndAlloc(void) {
330  if (delay_alloc) {
331  shandle = Storage::Get()->Alloc(shandle.size, shandle.ctx);
332  delay_alloc = false;
333  }
334  }
336  ~Chunk() {
337  if (static_data || delay_alloc) {
338  Engine::Get()->DeleteVariable([](RunContext s) {}, shandle.ctx, var);
339  } else {
340  Storage::Handle h = this->shandle;
341  Engine::Get()->DeleteVariable([h](RunContext s) {
342  Storage::Get()->Free(h);
343  }, shandle.ctx, var);
344  }
345  }
346  };
348  std::shared_ptr<Chunk> ptr_;
350  TShape shape_;
352  size_t offset_;
354  int dtype_;
355 };
356 
368 void CopyFromTo(const NDArray &from, NDArray *to, int priority = 0);
369 
376 void ElementwiseSum(const std::vector<NDArray> &source, NDArray *out, int priority = 0);
377 
384 NDArray operator+(const NDArray &lhs, const NDArray &rhs);
391 NDArray operator+(const NDArray &lhs, const real_t &rhs);
398 NDArray operator-(const NDArray &lhs, const NDArray &rhs);
405 NDArray operator-(const NDArray &lhs, const real_t &rhs);
412 NDArray operator*(const NDArray &lhs, const NDArray &rhs); \
419 NDArray operator*(const NDArray &lhs, const real_t &rhs);
426 NDArray operator/(const NDArray &lhs, const NDArray &rhs);
433 NDArray operator/(const NDArray &lhs, const real_t &rhs);
434 
439 void RandomSeed(uint32_t seed);
446 void SampleUniform(real_t begin, real_t end, NDArray *out);
447 
454 void SampleGaussian(real_t mu, real_t sigma, NDArray *out);
455 //--------------------------------------------------------------
456 // The following part are API Registration of NDArray functions.
457 //--------------------------------------------------------------
459 typedef std::function<void (NDArray **used_vars,
460  real_t *scalars,
461  NDArray **mutate_vars,
462  int num_params,
463  char **param_keys,
464  char **param_vals)> NDArrayAPIFunction;
480 };
483  : public dmlc::FunctionRegEntryBase<NDArrayFunctionReg,
484  NDArrayAPIFunction> {
486  unsigned num_use_vars;
488  unsigned num_mutate_vars;
490  unsigned num_scalars;
497  : num_use_vars(0),
498  num_mutate_vars(0),
499  num_scalars(0),
500  type_mask(0) {}
507  inline NDArrayFunctionReg &set_function(void (*fsetvalue)(const real_t &rhs,
508  NDArray *out)) {
509  body = [fsetvalue] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
510  int num_params, char **param_keys, char **param_vals) {
511  (*fsetvalue)(s[0], mutate_vars[0]);
512  };
513  num_mutate_vars = 1; num_scalars = 1;
514  this->add_argument("src", "real_t", "Source input to the function.");
515  return *this;
516  }
523  inline NDArrayFunctionReg &set_function(void(*fternary)(const NDArray &lhs,
524  const NDArray &mhs,
525  const NDArray &rhs,
526  NDArray *out)) {
527  body = [fternary](NDArray **used_vars,
528  real_t *s, NDArray **mutate_vars,
529  int num_params, char **param_keys, char **param_vals) {
530  (*fternary)(*used_vars[0], *used_vars[1], *used_vars[2], mutate_vars[0]);
531  };
532  num_use_vars = 3; num_mutate_vars = 1;
534  this->add_argument("lhs", "NDArray", "Left operand to the function.");
535  this->add_argument("mhs", "NDArray", "Middle operand to the function.");
536  this->add_argument("rhs", "NDArray", "Right operand to the function.");
537  return *this;
538  }
545  inline NDArrayFunctionReg &set_function(void (*fbinary)(const NDArray &lhs,
546  const NDArray &rhs,
547  NDArray *out)) {
548  body = [fbinary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
549  int num_params, char **param_keys, char **param_vals) {
550  (*fbinary)(*used_vars[0], *used_vars[1], mutate_vars[0]);
551  };
552  num_use_vars = 2; num_mutate_vars = 1;
554  this->add_argument("lhs", "NDArray", "Left operand to the function.");
555  this->add_argument("rhs", "NDArray", "Right operand to the function.");
556  return *this;
557  }
564  inline NDArrayFunctionReg &set_function(void (*fscalar)(const NDArray &lhs,
565  const real_t &rhs,
566  NDArray *out)) {
567  body = [fscalar] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
568  int num_params, char **param_keys, char **param_vals) {
569  (*fscalar)(*used_vars[0], s[0], mutate_vars[0]);
570  };
573  this->add_argument("lhs", "NDArray", "Left operand to the function.");
574  this->add_argument("rhs", "real_t", "Right operand to the function.");
575  return *this;
576  }
583  inline NDArrayFunctionReg &set_function(void (*funary)(const NDArray &src,
584  NDArray *out)) {
585  body = [funary] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
586  int num_params, char **param_keys, char **param_vals) {
587  (*funary)(*used_vars[0], mutate_vars[0]);
588  };
589  num_use_vars = 1; num_mutate_vars = 1;
591  this->add_argument("src", "NDArray", "Source input to the function.");
592  return *this;
593  }
601  void (*fgeneric)(NDArray **used_vars,
602  real_t *s,
603  NDArray **mutate_vars,
604  const std::map<std::string, std::string>& param)) {
605  body = [fgeneric] (NDArray **used_vars, real_t *s, NDArray **mutate_vars,
606  int num_params, char **param_keys, char **param_vals) {
607  std::map<std::string, std::string> param;
608  for (int i = 0; i < num_params; ++i) {
609  param[param_keys[i]] = param_vals[i];
610  }
611  fgeneric(used_vars, s, mutate_vars, param);
612  };
613  return *this;
614  }
620  inline NDArrayFunctionReg &set_num_use_vars(unsigned n) {
621  num_use_vars = n; return *this;
622  }
629  num_mutate_vars = n; return *this;
630  }
636  inline NDArrayFunctionReg &set_num_scalars(unsigned n) {
637  num_scalars = n; return *this;
638  }
644  inline NDArrayFunctionReg &set_type_mask(int tmask) {
645  type_mask = tmask; return *this;
646  }
647 }; // NDArrayFunctionReg
648 
660 #define MXNET_REGISTER_NDARRAY_FUN(name) \
661  DMLC_REGISTRY_REGISTER(::mxnet::NDArrayFunctionReg, NDArrayFunctionReg, name)
662 
663 } // namespace mxnet
664 
665 namespace dmlc {
667 DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true);
668 } // namespace dmlc
669 #endif // MXNET_NDARRAY_H_
bool Load(dmlc::Stream *strm)
load the content from binary stream
NDArrayFunctionReg & set_num_mutate_vars(unsigned n)
set the number of mutate variables
Definition: ndarray.h:628
Engine::VarHandle var() const
Definition: ndarray.h:112
void RandomSeed(uint32_t seed)
Seed the random number generator.
Engine that schedules all the operations according to dependency.
mshadow::TShape TShape
dynamic shape type
Definition: base.h:85
NDArray & operator/=(const NDArray &src)
elementwise division from current ndarray this mutate the current NDArray
void SyncCopyFromCPU(const void *data, size_t size) const
Do a synchronize copy from a continugous CPU memory region.
NDArrayFunctionReg()
constructor
Definition: ndarray.h:496
void PushSync(SyncFn exec_fn, Context exec_ctx, std::vector< VarHandle > const &const_vars, std::vector< VarHandle > const &mutable_vars, FnProperty prop=FnProperty::kNormal, int priority=0)
Push an synchronous operation to the engine.
Definition: engine.h:198
NDArray operator*(const NDArray &lhs, const NDArray &rhs)
elementwise multiplication
virtual void Free(Handle handle)=0
Free storage.
NDArray Slice(index_t begin, index_t end) const
Slice a NDArray.
Definition: ndarray.h:227
NDArrayFunctionReg & set_num_use_vars(unsigned n)
set the number of mutate variables
Definition: ndarray.h:620
static Context CPU()
DMLC_DECLARE_TRAITS(has_saveload, mxnet::NDArray, true)
traits
mshadow::default_real_t real_t
data type that will be used to store ndarray
Definition: base.h:82
int type_mask
information on how function should be called from API
Definition: ndarray.h:492
NDArray T() const
return transpose of current NDArray
NDArrayFunctionReg & set_function(void(*funary)(const NDArray &src, NDArray *out))
set the function body to a unary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:583
NDArrayFunctionReg & set_num_scalars(unsigned n)
set the number of scalar arguments
Definition: ndarray.h:636
mshadow::TBlob TBlob
storage container type
Definition: base.h:87
NDArray & operator+=(const NDArray &src)
elementwise add to current space this mutate the current NDArray
unsigned num_mutate_vars
number of variable mutated by this function
Definition: ndarray.h:488
execution time context. The information needed in runtime for actual execution.
Definition: base.h:181
NDArrayFunctionReg & set_function(void(*fscalar)(const NDArray &lhs, const real_t &rhs, NDArray *out))
set the function body to a binary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:564
void SyncCopyToCPU(void *data, size_t size) const
Do a synchronize copy to a continugous CPU memory region.
NDArray()
default cosntructor
Definition: ndarray.h:34
unsigned num_use_vars
number of variable used by this function
Definition: ndarray.h:486
NDArrayFunctionReg & set_function(void(*fternary)(const NDArray &lhs, const NDArray &mhs, const NDArray &rhs, NDArray *out))
set the function body to a ternary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:523
virtual Handle Alloc(size_t size, Context ctx)=0
Allocate a new contiguous memory for a given size.
NDArray Reshape(const TShape &shape) const
Get an reshaped NDArray.
Definition: ndarray.h:255
bool is_none() const
Definition: ndarray.h:87
all the scalar should go before use_vars
Definition: ndarray.h:470
virtual VarHandle NewVariable()=0
Allocate a new variable, the variable can then be used to schedule the operation concurrently via dep...
NDArray & operator=(real_t scalar)
set all the elements in ndarray to be scalar
whether this function allows the handles in the target to be empty NDArray that are not yet initializ...
Definition: ndarray.h:479
static Storage * Get()
const TShape & shape() const
Definition: ndarray.h:61
virtual void WaitForVar(VarHandle var)=0
Wait for a variable.
NDArray & operator-=(const NDArray &src)
elementwise subtract from current ndarray this mutate the current NDArray
Context ctx() const
Definition: ndarray.h:77
void SampleGaussian(real_t mu, real_t sigma, NDArray *out)
Sample gaussian distribution for each elements of out.
Storage manager across multiple devices.
void WaitToRead() const
Block until all the pending write operations with respect to current NDArray are finished, and read can be performed.
Definition: ndarray.h:94
int dtype() const
Definition: ndarray.h:83
Storage handle.
Definition: storage.h:22
virtual void DeleteVariable(SyncFn delete_fn, Context exec_ctx, VarHandle var)=0
Schedule the deletion of a variable.
NDArrayFunctionReg & set_type_mask(int tmask)
set type mask
Definition: ndarray.h:644
TBlob data() const
Definition: ndarray.h:67
engine::VarHandle VarHandle
Variable pointer.
Definition: engine.h:83
void WaitToWrite() const
Block until all the pending read/write operations with respect to current NDArray are finished...
Definition: ndarray.h:102
NDArray operator-(const NDArray &lhs, const NDArray &rhs)
elementwise substraction
NDArrayFunctionReg & set_function(void(*fsetvalue)(const real_t &rhs, NDArray *out))
set the function body to a NDArray setvalue function this will also auto set the parameters correctly...
Definition: ndarray.h:507
NDArray operator+(const NDArray &lhs, const NDArray &rhs)
elementwise add
void SampleUniform(real_t begin, real_t end, NDArray *out)
Sample uniform distribution for each elements of out.
Registry entry for NDArrayFunction.
Definition: ndarray.h:482
NDArray At(index_t idx) const
Index a NDArray.
Definition: ndarray.h:241
NDArrayFunctionReg & set_function(void(*fbinary)(const NDArray &lhs, const NDArray &rhs, NDArray *out))
set the function body to a binary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:545
NDArray Copy(Context ctx) const
return a new copy this NDArray
configuation of mxnet as well as basic data structure.
NDArray & operator*=(const NDArray &src)
elementwise multiplication to current ndarray this mutate the current NDArray
all the use_vars should go before scalar
Definition: ndarray.h:468
unsigned num_scalars
number of scalars used by this function
Definition: ndarray.h:490
static Engine * Get()
void Save(dmlc::Stream *strm) const
save the content into binary stream
void CheckAndAlloc() const
Allocate the space if it is delayed allocated. This is an internal function used by system that norma...
Definition: ndarray.h:266
mshadow::index_t index_t
index type usually use unsigned
Definition: base.h:80
static Context GPU(int32_t dev_id)
Context information about the execution enviroment.
Definition: base.h:90
ndarray interface
Definition: ndarray.h:31
void CopyFromTo(const NDArray &from, NDArray *to, int priority=0)
issue an copy operation from one NDArray to another the two ndarray can sit on different devices this...
NDArray(const TBlob &data, int dev_id)
constructing a static NDArray that shares data with TBlob Use with caution: allocate ONLY ONE NDArray...
Definition: ndarray.h:54
void ElementwiseSum(const std::vector< NDArray > &source, NDArray *out, int priority=0)
Perform elementwise sum over each data from source, store result into out.
std::function< void(NDArray **used_vars, real_t *scalars, NDArray **mutate_vars, int num_params, char **param_keys, char **param_vals)> NDArrayAPIFunction
definition of NDArray function
Definition: ndarray.h:464
NDArrayFunctionReg & set_function(void(*fgeneric)(NDArray **used_vars, real_t *s, NDArray **mutate_vars, const std::map< std::string, std::string > &param))
set the function body to a unary NDArray function this will also auto set the parameters correctly ...
Definition: ndarray.h:600
NDArray(const TShape &shape, Context ctx, bool delay_alloc=false, int dtype=mshadow::default_type_flag)
constructing a new dynamic NDArray
Definition: ndarray.h:42
NDArray operator/(const NDArray &lhs, const NDArray &rhs)
elementwise division
NDArrayFunctionTypeMask
mask information on how functions can be exposed
Definition: ndarray.h:466