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,
          int dtype        = 0);
  NDArray(const Shape& shape, const Context& context, bool delay_alloc = true, int dtype = 0);
  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 index) 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 LoadFromBuffer(const void* buffer,
                             size_t size,
                             std::vector<NDArray>* array_list          = nullptr,
                             std::map<std::string, NDArray>* array_map = nullptr);
  static std::map<std::string, NDArray> LoadFromBufferToMap(const void* buffer, size_t size);
  static std::vector<NDArray> LoadFromBufferToList(const void* buffer, size_t size);
  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_