forked from mindspore-Ecosystem/mindspore
Provide Tensor.from_numpy() to avoid copy
This commit is contained in:
parent
c9d584fa74
commit
d7cdda1cb0
|
@ -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.
|
||||
|
||||
|
|
|
@ -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);
|
||||
|
|
|
@ -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;
|
||||
|
|
|
@ -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;
|
||||
};
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue