Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
467 lines (443 sloc) 15 KB
/*
* 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.
*/
/*!
* \file ndarray.hpp
* \brief implementation of the ndarray
* \author Zhang Chen, Chuntao Hong
*/
#ifndef MXNET_CPP_NDARRAY_HPP_
#define MXNET_CPP_NDARRAY_HPP_
#include <algorithm>
#include <map>
#include <string>
#include <vector>
#include <iterator>
#include "dmlc/logging.h"
#include "mxnet-cpp/ndarray.h"
#include "mxnet-cpp/operator.h"
namespace mxnet {
namespace cpp {
inline NDArray::NDArray() {
NDArrayHandle handle;
CHECK_EQ(MXNDArrayCreateNone(&handle), 0);
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray::NDArray(const NDArrayHandle &handle) {
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray::NDArray(const std::vector<mx_uint> &shape, const Context &context,
bool delay_alloc, int dtype) {
NDArrayHandle handle;
CHECK_EQ(MXNDArrayCreateEx(shape.data(), shape.size(), context.GetDeviceType(),
context.GetDeviceId(), delay_alloc, dtype, &handle),
0);
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray::NDArray(const Shape &shape, const Context &context,
bool delay_alloc, int dtype) {
NDArrayHandle handle;
CHECK_EQ(MXNDArrayCreateEx(shape.data(), shape.ndim(), context.GetDeviceType(),
context.GetDeviceId(), delay_alloc, dtype, &handle),
0);
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray::NDArray(const mx_float *data, size_t size) {
NDArrayHandle handle;
CHECK_EQ(MXNDArrayCreateNone(&handle), 0);
MXNDArraySyncCopyFromCPU(handle, data, size);
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray::NDArray(const mx_float *data, const Shape &shape,
const Context &context) {
NDArrayHandle handle;
CHECK_EQ(MXNDArrayCreate(shape.data(), shape.ndim(), context.GetDeviceType(),
context.GetDeviceId(), false, &handle),
0);
MXNDArraySyncCopyFromCPU(handle, data, shape.Size());
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray::NDArray(const std::vector<mx_float> &data, const Shape &shape,
const Context &context) {
NDArrayHandle handle;
CHECK_EQ(MXNDArrayCreate(shape.data(), shape.ndim(), context.GetDeviceType(),
context.GetDeviceId(), false, &handle),
0);
MXNDArraySyncCopyFromCPU(handle, data.data(), shape.Size());
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray::NDArray(const std::vector<mx_float> &data) {
NDArrayHandle handle;
CHECK_EQ(MXNDArrayCreateNone(&handle), 0);
MXNDArraySyncCopyFromCPU(handle, data.data(), data.size());
blob_ptr_ = std::make_shared<NDBlob>(handle);
}
inline NDArray NDArray::operator+(mx_float scalar) {
NDArray ret;
Operator("_plus_scalar")(*this, scalar).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator-(mx_float scalar) {
NDArray ret;
Operator("_minus_scalar")(*this, scalar).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator*(mx_float scalar) {
NDArray ret;
Operator("_mul_scalar")(*this, scalar).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator/(mx_float scalar) {
NDArray ret;
Operator("_div_scalar")(*this, scalar).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator%(mx_float scalar) {
NDArray ret;
Operator("_mod_scalar")(*this, scalar).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator+(const NDArray &rhs) {
NDArray ret;
Operator("_plus")(*this, rhs).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator-(const NDArray &rhs) {
NDArray ret;
Operator("_minus")(*this, rhs).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator*(const NDArray &rhs) {
NDArray ret;
Operator("_mul")(*this, rhs).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator/(const NDArray &rhs) {
NDArray ret;
Operator("_div")(*this, rhs).Invoke(ret);
return ret;
}
inline NDArray NDArray::operator%(const NDArray &rhs) {
NDArray ret;
Operator("_mod")(*this, rhs).Invoke(ret);
return ret;
}
inline NDArray &NDArray::operator=(mx_float scalar) {
Operator("_set_value")(scalar).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator+=(mx_float scalar) {
Operator("_plus_scalar")(*this, scalar).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator-=(mx_float scalar) {
Operator("_minus_scalar")(*this, scalar).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator*=(mx_float scalar) {
Operator("_mul_scalar")(*this, scalar).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator/=(mx_float scalar) {
Operator("_div_scalar")(*this, scalar).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator%=(mx_float scalar) {
Operator("_mod_scalar")(*this, scalar).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator+=(const NDArray &rhs) {
Operator("_plus")(*this, rhs).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator-=(const NDArray &rhs) {
Operator("_minus")(*this, rhs).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator*=(const NDArray &rhs) {
Operator("_mul")(*this, rhs).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator/=(const NDArray &rhs) {
Operator("_div")(*this, rhs).Invoke(*this);
return *this;
}
inline NDArray &NDArray::operator%=(const NDArray &rhs) {
Operator("_mod")(*this, rhs).Invoke(*this);
return *this;
}
inline NDArray NDArray::ArgmaxChannel() {
NDArray ret;
Operator("argmax_channel")(*this).Invoke(ret);
return ret;
}
inline void NDArray::SyncCopyFromCPU(const mx_float *data, size_t size) {
MXNDArraySyncCopyFromCPU(blob_ptr_->handle_, data, size);
}
inline void NDArray::SyncCopyFromCPU(const std::vector<mx_float> &data) {
MXNDArraySyncCopyFromCPU(blob_ptr_->handle_, data.data(), data.size());
}
inline void NDArray::SyncCopyToCPU(mx_float *data, size_t size) {
MXNDArraySyncCopyToCPU(blob_ptr_->handle_, data, size > 0 ? size : Size());
}
inline void NDArray::SyncCopyToCPU(std::vector<mx_float> *data, size_t size) {
size = size > 0 ? size : Size();
data->resize(size);
MXNDArraySyncCopyToCPU(blob_ptr_->handle_, data->data(), size);
}
inline NDArray NDArray::Copy(const Context &ctx) const {
NDArray ret(GetShape(), ctx, true, this->GetDType());
Operator("_copyto")(*this).Invoke(ret);
return ret;
}
inline NDArray NDArray::CopyTo(NDArray * other) const {
Operator("_copyto")(*this).Invoke(*other);
return *other;
}
inline NDArray NDArray::Slice(mx_uint begin, mx_uint end) const {
NDArrayHandle handle;
CHECK_EQ(MXNDArraySlice(GetHandle(), begin, end, &handle), 0);
return NDArray(handle);
}
inline NDArray NDArray::Reshape(const Shape &new_shape) const {
NDArrayHandle handle;
std::vector<int> dims(new_shape.ndim());
for (index_t i = 0; i < new_shape.ndim(); ++i) {
dims[i] = new_shape[i];
}
new_shape.data();
CHECK_EQ(
MXNDArrayReshape(GetHandle(), new_shape.ndim(), dims.data(), &handle), 0);
return NDArray(handle);
}
inline void NDArray::WaitToRead() const {
CHECK_EQ(MXNDArrayWaitToRead(blob_ptr_->handle_), 0) << MXGetLastError();
}
inline void NDArray::WaitToWrite() {
CHECK_EQ(MXNDArrayWaitToWrite(blob_ptr_->handle_), 0) << MXGetLastError();
}
inline void NDArray::WaitAll() { CHECK_EQ(MXNDArrayWaitAll(), 0) << MXGetLastError(); }
inline void NDArray::SampleGaussian(mx_float mu, mx_float sigma, NDArray *out) {
Operator("_random_normal")(mu, sigma).Invoke(*out);
}
inline void NDArray::SampleUniform(mx_float begin, mx_float end, NDArray *out) {
Operator("_random_uniform")(begin, end).Invoke(*out);
}
inline void NDArray::Load(const std::string &file_name,
std::vector<NDArray> *array_list,
std::map<std::string, NDArray> *array_map) {
mx_uint out_size, out_name_size;
NDArrayHandle *out_arr;
const char **out_names;
CHECK_EQ(MXNDArrayLoad(file_name.c_str(), &out_size, &out_arr, &out_name_size,
&out_names),
0);
if (array_list != nullptr) {
array_list->reserve(out_size);
for (mx_uint i = 0; i < out_size; ++i) {
array_list->push_back(NDArray(out_arr[i]));
}
}
if (array_map != nullptr && out_name_size > 0) {
CHECK_EQ(out_name_size, out_size);
for (mx_uint i = 0; i < out_size; ++i) {
(*array_map)[out_names[i]] = NDArray(out_arr[i]);
}
}
}
inline std::map<std::string, NDArray> NDArray::LoadToMap(
const std::string &file_name) {
std::map<std::string, NDArray> array_map;
mx_uint out_size, out_name_size;
NDArrayHandle *out_arr;
const char **out_names;
CHECK_EQ(MXNDArrayLoad(file_name.c_str(), &out_size, &out_arr, &out_name_size,
&out_names),
0);
if (out_name_size > 0) {
CHECK_EQ(out_name_size, out_size);
for (mx_uint i = 0; i < out_size; ++i) {
array_map[out_names[i]] = NDArray(out_arr[i]);
}
}
return array_map;
}
inline std::vector<NDArray> NDArray::LoadToList(const std::string &file_name) {
std::vector<NDArray> array_list;
mx_uint out_size, out_name_size;
NDArrayHandle *out_arr;
const char **out_names;
CHECK_EQ(MXNDArrayLoad(file_name.c_str(), &out_size, &out_arr, &out_name_size,
&out_names),
0);
array_list.reserve(out_size);
for (mx_uint i = 0; i < out_size; ++i) {
array_list.push_back(NDArray(out_arr[i]));
}
return array_list;
}
inline void NDArray::LoadFromBuffer(const void *buffer, size_t size,
std::vector<NDArray> *array_list,
std::map<std::string, NDArray> *array_map) {
mx_uint out_size, out_name_size;
NDArrayHandle *out_arr;
const char **out_names;
CHECK_EQ(MXNDArrayLoadFromBuffer(buffer, size, &out_size, &out_arr, &out_name_size,
&out_names),
0);
if (array_list != nullptr) {
array_list->reserve(out_size);
for (mx_uint i = 0; i < out_size; ++i) {
array_list->push_back(NDArray(out_arr[i]));
}
}
if (array_map != nullptr && out_name_size > 0) {
CHECK_EQ(out_name_size, out_size);
for (mx_uint i = 0; i < out_size; ++i) {
(*array_map)[out_names[i]] = NDArray(out_arr[i]);
}
}
}
inline std::map<std::string, NDArray> NDArray::LoadFromBufferToMap(
const void *buffer, size_t size) {
std::map<std::string, NDArray> array_map;
mx_uint out_size, out_name_size;
NDArrayHandle *out_arr;
const char **out_names;
CHECK_EQ(MXNDArrayLoadFromBuffer(buffer, size, &out_size, &out_arr, &out_name_size,
&out_names),
0);
if (out_name_size > 0) {
CHECK_EQ(out_name_size, out_size);
for (mx_uint i = 0; i < out_size; ++i) {
array_map[out_names[i]] = NDArray(out_arr[i]);
}
}
return array_map;
}
inline std::vector<NDArray> NDArray::LoadFromBufferToList(const void *buffer, size_t size) {
std::vector<NDArray> array_list;
mx_uint out_size, out_name_size;
NDArrayHandle *out_arr;
const char **out_names;
CHECK_EQ(MXNDArrayLoadFromBuffer(buffer, size, &out_size, &out_arr, &out_name_size,
&out_names),
0);
array_list.reserve(out_size);
for (mx_uint i = 0; i < out_size; ++i) {
array_list.push_back(NDArray(out_arr[i]));
}
return array_list;
}
inline void NDArray::Save(const std::string &file_name,
const std::map<std::string, NDArray> &array_map) {
std::vector<NDArrayHandle> args;
std::vector<const char *> keys;
for (const auto &t : array_map) {
args.push_back(t.second.GetHandle());
keys.push_back(t.first.c_str());
}
CHECK_EQ(
MXNDArraySave(file_name.c_str(), args.size(), args.data(), keys.data()),
0);
}
inline void NDArray::Save(const std::string &file_name,
const std::vector<NDArray> &array_list) {
std::vector<NDArrayHandle> args;
for (const auto &t : array_list) {
args.push_back(t.GetHandle());
}
CHECK_EQ(MXNDArraySave(file_name.c_str(), args.size(), args.data(), nullptr),
0);
}
inline size_t NDArray::Offset(size_t h, size_t w) const {
auto const shape = GetShape();
CHECK_EQ(shape.size(), 2) << "The NDArray needs to be 2 dimensional.";
return (h * shape[1]) + w;
}
inline size_t NDArray::Offset(size_t c, size_t h, size_t w) const {
auto const shape = GetShape();
CHECK_EQ(shape.size(), 3) << "The NDArray needs to be 3 dimensional.";
return h * shape[0] * shape[2] + w * shape[0] + c;
}
inline mx_float NDArray::At(size_t h, size_t w) const {
return GetData()[Offset(h, w)];
}
inline mx_float NDArray::At(size_t c, size_t h, size_t w) const {
return GetData()[Offset(c, h, w)];
}
inline mx_float NDArray::At(size_t index) const {
auto shape = GetShape();
CHECK_EQ(shape.size(), 1) << "The NDArray needs to be 1 dimensional.";
CHECK_LT(index, shape[0]) << "Specified index is out of range.";
return GetData()[index];
}
inline size_t NDArray::Size() const {
size_t ret = 1;
for (auto &i : GetShape()) ret *= i;
return ret;
}
inline std::vector<mx_uint> NDArray::GetShape() const {
const int *out_pdata;
int out_dim;
MXNDArrayGetShapeEx(blob_ptr_->handle_, &out_dim, &out_pdata);
std::vector<mx_uint> ret;
for (int i = 0; i < out_dim; ++i) {
ret.push_back(out_pdata[i]);
}
return ret;
}
inline int NDArray::GetDType() const {
int ret;
MXNDArrayGetDType(blob_ptr_->handle_, &ret);
return ret;
}
inline const mx_float *NDArray::GetData() const {
void *ret;
MXNDArrayGetData(blob_ptr_->handle_, &ret);
if (GetDType() != 0) {
return NULL;
}
return static_cast<mx_float*>(ret);
}
inline Context NDArray::GetContext() const {
int out_dev_type;
int out_dev_id;
MXNDArrayGetContext(blob_ptr_->handle_, &out_dev_type, &out_dev_id);
return Context((DeviceType)out_dev_type, out_dev_id);
}
inline std::ostream & operator<<(std::ostream &out, const NDArray &ndarray) {
// TODO(lx75249): Consider DType / beautify like numpy
auto shape = ndarray.GetShape();
NDArray cpu_array(ndarray.GetShape(), Context::cpu());
if (ndarray.GetContext().GetDeviceType() != DeviceType::kGPU) {
cpu_array = ndarray;
} else {
ndarray.WaitToRead();
ndarray.CopyTo(&cpu_array);
}
out << '[';
cpu_array.WaitToRead();
std::copy(cpu_array.GetData(), cpu_array.GetData() + ndarray.Size(),
std::ostream_iterator<float>(out, ", "));
out << ']';
return out;
}
} // namespace cpp
} // namespace mxnet
#endif // MXNET_CPP_NDARRAY_HPP_
You can’t perform that action at this time.