ndarray.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

#ifndef MXNET_CPP_NDARRAY_H_
#define MXNET_CPP_NDARRAY_H_

#include <map>
#include <memory>
#include <string>
#include <vector>
#include <iostream>
#include "mxnet-cpp/base.h"
#include "mxnet-cpp/shape.h"

namespace mxnet {
namespace cpp {

enum DeviceType {
  kCPU = 1,
  kGPU = 2,
  kCPUPinned = 3
};

class Context {
 public:
  Context(const DeviceType &type, int id) : type_(type), id_(id) {}
  DeviceType GetDeviceType() const { return type_; }
  int GetDeviceId() const { return id_; }

  static Context gpu(int device_id = 0) {
    return Context(DeviceType::kGPU, device_id);
  }

  static Context cpu(int device_id = 0) {
    return Context(DeviceType::kCPU, device_id);
  }

 private:
  DeviceType type_;
  int id_;
};

struct NDBlob {
 public:
  NDBlob() : handle_(nullptr) {}
  explicit NDBlob(NDArrayHandle handle) : handle_(handle) {}
  ~NDBlob() { MXNDArrayFree(handle_); }
  NDArrayHandle handle_;

 private:
  NDBlob(const NDBlob &);
  NDBlob &operator=(const NDBlob &);
};

class NDArray {
 public:
  NDArray();
  explicit NDArray(const NDArrayHandle &handle);
  NDArray(const std::vector<mx_uint> &shape, const Context &context,
          bool delay_alloc = true);
  NDArray(const Shape &shape, const Context &context, bool delay_alloc = true);
  NDArray(const mx_float *data, size_t size);
  NDArray(const mx_float *data, const Shape &shape, const Context &context);
  NDArray(const std::vector<mx_float> &data, const Shape &shape,
          const Context &context);
  explicit NDArray(const std::vector<mx_float> &data);
  NDArray operator+(mx_float scalar);
  NDArray operator-(mx_float scalar);
  NDArray operator*(mx_float scalar);
  NDArray operator/(mx_float scalar);
  NDArray operator%(mx_float scalar);
  NDArray operator+(const NDArray &);
  NDArray operator-(const NDArray &);
  NDArray operator*(const NDArray &);
  NDArray operator/(const NDArray &);
  NDArray operator%(const NDArray &);
  NDArray &operator=(mx_float scalar);
  NDArray &operator+=(mx_float scalar);
  NDArray &operator-=(mx_float scalar);
  NDArray &operator*=(mx_float scalar);
  NDArray &operator/=(mx_float scalar);
  NDArray &operator%=(mx_float scalar);
  NDArray &operator+=(const NDArray &src);
  NDArray &operator-=(const NDArray &src);
  NDArray &operator*=(const NDArray &src);
  NDArray &operator/=(const NDArray &src);
  NDArray &operator%=(const NDArray &src);
  NDArray ArgmaxChannel();
  void SyncCopyFromCPU(const mx_float *data, size_t size);
  void SyncCopyFromCPU(const std::vector<mx_float> &data);
  void SyncCopyToCPU(mx_float *data, size_t size = 0);
  void SyncCopyToCPU(std::vector<mx_float> *data, size_t size = 0);
  NDArray CopyTo(NDArray * other) const;
  NDArray Copy(const Context &) const;
  size_t Offset(size_t h = 0, size_t w = 0) const;
  size_t Offset(size_t c, size_t h, size_t w) const;
  mx_float At(size_t h, size_t w) const;
  mx_float At(size_t c, size_t h, size_t w) const;
  NDArray Slice(mx_uint begin, mx_uint end) const;
  NDArray Reshape(const Shape &new_shape) const;
  void WaitToRead() const;
  void WaitToWrite();
  static void WaitAll();
  static void SampleGaussian(mx_float mu, mx_float sigma, NDArray *out);
  static void SampleUniform(mx_float begin, mx_float end, NDArray *out);
  static void Load(const std::string &file_name,
                   std::vector<NDArray> *array_list = nullptr,
                   std::map<std::string, NDArray> *array_map = nullptr);
  static std::map<std::string, NDArray> LoadToMap(const std::string &file_name);
  static std::vector<NDArray> LoadToList(const std::string &file_name);
  static void Save(const std::string &file_name,
                   const std::map<std::string, NDArray> &array_map);
  static void Save(const std::string &file_name,
                   const std::vector<NDArray> &array_list);
  size_t Size() const;
  std::vector<mx_uint> GetShape() const;
  int GetDType() const;
  const mx_float *GetData() const;

  Context GetContext() const;

  NDArrayHandle GetHandle() const { return blob_ptr_->handle_; }

 private:
  std::shared_ptr<NDBlob> blob_ptr_;
};

std::ostream& operator<<(std::ostream& out, const NDArray &ndarray);
}  // namespace cpp
}  // namespace mxnet

#endif  // MXNET_CPP_NDARRAY_H_