Provide Tensor.from_numpy() to avoid copy

This commit is contained in:
He Wei 2020-08-27 20:58:44 +08:00
parent c9d584fa74
commit d7cdda1cb0
5 changed files with 152 additions and 10 deletions

View File

@ -117,6 +117,76 @@ static bool IsCContiguous(const py::array &input) {
return (flags & pybind11::detail::npy_api::NPY_ARRAY_C_CONTIGUOUS_) != 0;
}
// TensorDataNumpy implements TensorData using numpy array.
class TensorDataNumpy : public TensorData {
public:
explicit TensorDataNumpy(const py::array &input) : data_(input) {
if (!IsCContiguous(data_)) {
// Call numpy.ascontiguousarray() to convert data to C contiguous if it is not.
auto np = py::module::import("numpy");
auto convert = np.attr("ascontiguousarray");
data_ = convert(data_);
}
}
/// Total number of elements.
ssize_t size() const override { return data_.size(); }
/// Byte size of a single element.
ssize_t itemsize() const override { return data_.itemsize(); }
/// Total number of bytes.
ssize_t nbytes() const override { return data_.nbytes(); }
/// Number of dimensions.
ssize_t ndim() const override { return data_.ndim(); }
/// Data pointer.
void *data() override { return data_.request().ptr; }
const void *const_data() const override { return data_.request().ptr; }
/// Is data equals.
bool equals(const TensorData &other) const override {
auto ptr = dynamic_cast<const TensorDataNumpy *>(&other);
if (ptr == nullptr) {
// Not same type, compare data byte by byte.
return TensorData::equals(other);
}
return NumpyEquals(*ptr);
}
bool NumpyEquals(const TensorDataNumpy &other) const {
auto all_data_equal = [&other, this]() -> bool {
auto np = py::module::import("numpy");
auto equal = np.attr("equal")(data_, other.data_);
auto all_equal = np.attr("all")(equal);
return all_equal.cast<bool>();
};
return this == &other || data_.is(other.data_) || all_data_equal();
}
/// To string.
std::string ToString(const TypeId type, const ShapeVector &shape, bool use_comma) const override {
if (use_comma) {
// Call python np.array2string(data_, separator=', ') to convert string with comma.
py::dict kwargs;
kwargs["separator"] = ", ";
auto np = py::module::import("numpy");
auto array2string = np.attr("array2string");
return py::str(array2string(data_, **kwargs));
}
// without comma.
return py::str(data_);
}
/// py::array object.
py::array py_array() const { return data_; }
private:
mutable py::array data_;
};
TensorPtr TensorPy::MakeTensor(const py::array &input, const TypePtr &type_ptr) {
// Get input buffer info.
py::buffer_info buf = input.request();
@ -145,7 +215,7 @@ TensorPtr TensorPy::MakeTensor(const py::array &input, const TypePtr &type_ptr)
buf.ptr = tmp_buf.get();
}
// Get tensor shape.
std::vector<int> shape(buf.shape.begin(), buf.shape.end());
ShapeVector shape(buf.shape.begin(), buf.shape.end());
if (data_type == buf_type) {
// Use memory copy if input data type is the same as the required type.
return std::make_shared<Tensor>(data_type, shape, buf.ptr, buf.size * buf.itemsize);
@ -154,6 +224,22 @@ TensorPtr TensorPy::MakeTensor(const py::array &input, const TypePtr &type_ptr)
return std::make_shared<Tensor>(data_type, shape, buf.ptr, buf_type);
}
/// Creates a Tensor from a numpy array without copy
TensorPtr TensorPy::MakeTensorNoCopy(const py::array &input) {
// Get input buffer info.
py::buffer_info buf = input.request();
// Get tensor dtype and check it.
auto dtype = GetDataType(buf);
if (dtype == TypeId::kTypeUnknown) {
MS_LOG(EXCEPTION) << "Unsupported data type!";
}
// Get tensor shape.
ShapeVector shape(buf.shape.begin(), buf.shape.end());
// Make a tensor with shared data with numpy array.
auto tensor_data = std::make_shared<TensorDataNumpy>(input);
return std::make_shared<Tensor>(dtype, shape, tensor_data);
}
static std::vector<ssize_t> GetStrides(const std::vector<ssize_t> &shape, ssize_t item_size) {
std::vector<ssize_t> strides;
strides.reserve(shape.size());
@ -186,19 +272,23 @@ py::tuple TensorPy::GetPyTupleShape(const Tensor &tensor) {
py::array TensorPy::SyncAsNumpy(const Tensor &tensor) {
tensor.data_sync();
auto info = GetPyBufferInfo(tensor);
py::object self = py::cast(&tensor);
return py::array(py::dtype(info), info.shape, info.strides, info.ptr, self);
return AsNumpy(tensor);
}
py::array TensorPy::AsNumpy(const Tensor &tensor) {
auto data_numpy = dynamic_cast<const TensorDataNumpy *>(&tensor.data());
if (data_numpy) {
// Return internal numpy array if tensor data is implemented base on it.
return data_numpy->py_array();
}
// Otherwise, create numpy array by buffer protocol.
auto info = GetPyBufferInfo(tensor);
py::object self = py::cast(&tensor);
return py::array(py::dtype(info), info.shape, info.strides, info.ptr, self);
}
static std::vector<int> GetShapeFromTuple(const py::tuple &tuple) {
std::vector<int> shape;
static ShapeVector GetShapeFromTuple(const py::tuple &tuple) {
ShapeVector shape;
const size_t size = tuple.size();
shape.reserve(tuple.size());
for (size_t i = 0; i < size; ++i) {
@ -210,7 +300,7 @@ static std::vector<int> GetShapeFromTuple(const py::tuple &tuple) {
REGISTER_PYBIND_DEFINE(Tensor, ([](const py::module *m) {
// Define python MetaTensor class.
(void)py::class_<MetaTensor, std::shared_ptr<MetaTensor>>(*m, "MetaTensor")
.def(py::init<TypePtr, const std::vector<int>>(), py::arg("dtype"), py::arg("shape"))
.def(py::init<TypePtr, const ShapeVector>(), py::arg("dtype"), py::arg("shape"))
.def_property_readonly("dtype", &MetaTensor::Dtype, "Get the MetaTensor's dtype.")
.def_property_readonly("shape", &MetaTensor::shape, "Get the MetaTensor's shape.")
.def_property("_param_info", &MetaTensor::param_info, &MetaTensor::set_param_info)
@ -224,7 +314,7 @@ REGISTER_PYBIND_DEFINE(Tensor, ([](const py::module *m) {
throw std::runtime_error("Invalid state!");
}
/* Create a new C++ instance */
MetaTensor tensor(TypeId(t[0].cast<int>()), t[1].cast<std::vector<int>>());
MetaTensor tensor(TypeId(t[0].cast<int>()), t[1].cast<ShapeVector>());
return tensor;
}));
// Define python Tensor class.
@ -288,6 +378,19 @@ REGISTER_PYBIND_DEFINE(Tensor, ([](const py::module *m) {
>>> data.shape()
(3, 3)
)mydelimiter")
.def("from_numpy", TensorPy::MakeTensorNoCopy, R"mydelimiter(
Creates a Tensor from a numpy.ndarray without copy.
Arg:
array (numpy.ndarray): The input ndarray.
Returns:
Tensor, tensor with shared data to input ndarray.
Examples:
>>> a = np.ones((2, 3))
>>> t = mindspore.Tensor.from_numpy(a)
)mydelimiter")
.def("asnumpy", TensorPy::SyncAsNumpy, R"mydelimiter(
Convert tensor to numpy.ndarray.

View File

@ -99,6 +99,11 @@ class TensorPy {
// param data_type [TypeId] Data type of the tensor.
static TensorPtr MakeTensor(const py::array &input, const TypePtr &data_type = nullptr);
// brief Create Tensor from a numpy array without copy.
//
// param input [py::array] Data value of the tensor.
static TensorPtr MakeTensorNoCopy(const py::array &input);
static py::array SyncAsNumpy(const Tensor &tensor);
static py::array AsNumpy(const Tensor &tensor);

View File

@ -198,10 +198,16 @@ class TensorDataImpl : public TensorData {
return data_.get();
}
const void *const_data() const override {
// May return nullptr if data not initialized.
return data_.get();
}
bool equals(const TensorData &other) const override {
auto ptr = dynamic_cast<const TensorDataImpl<T> *>(&other);
if (ptr == nullptr) {
return false;
// Not same type, compare data byte by byte.
return TensorData::equals(other);
}
if (ptr == this) {
return true;

View File

@ -50,8 +50,23 @@ class TensorData {
virtual ssize_t ndim() const = 0;
/// Data pointer.
virtual void *data() = 0;
/// Const Data pointer.
virtual const void *const_data() const = 0;
/// Is data equals.
virtual bool equals(const TensorData &other) const = 0;
virtual bool equals(const TensorData &other) const {
if (this == &other) {
return true;
}
// By default, compare data byte by byte.
auto this_data = static_cast<const uint8_t *>(const_data());
auto other_data = static_cast<const uint8_t *>(other.const_data());
if (this_data == nullptr || other_data == nullptr) {
// null means data not initialized, compare uninitialized data always return false.
return false;
}
return (this_data == other_data) || (ndim() == other.ndim() && nbytes() == other.nbytes() &&
std::equal(this_data, this_data + nbytes(), other_data));
}
/// To string.
virtual std::string ToString(const TypeId type, const ShapeVector &shape, bool use_comma) const = 0;
};

View File

@ -476,3 +476,16 @@ def test_tensor_operation():
assert np.all(x.asnumpy() == np.ones((3, 3)))
res = 5 // x
assert np.all(x.asnumpy() == np.ones((3, 3)))
def test_tensor_from_numpy():
a = np.ones((2, 3))
t = ms.Tensor.from_numpy(a)
assert np.all(t.asnumpy() == 1)
# 't' and 'a' share same data.
a[1] = 2
assert np.all(t.asnumpy()[0] == 1)
assert np.all(t.asnumpy()[1] == 2)
# 't' is still valid after 'a' deleted.
del a
assert np.all(t.asnumpy()[0] == 1)
assert np.all(t.asnumpy()[1] == 2)