From cfbc10a49147169b8c3b4aa128ee9f1ddd7d6c0a Mon Sep 17 00:00:00 2001 From: wyg1997 Date: Thu, 16 Dec 2021 18:41:45 +0800 Subject: [PATCH] Fix(FromNumpy): fix bug in stride --- oneflow/api/python/functional/tensor_api.cpp | 15 +++++++++++++-- python/oneflow/test/modules/test_from_numpy.py | 3 +++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/oneflow/api/python/functional/tensor_api.cpp b/oneflow/api/python/functional/tensor_api.cpp index 10c162dd182..ca24543062e 100644 --- a/oneflow/api/python/functional/tensor_api.cpp +++ b/oneflow/api/python/functional/tensor_api.cpp @@ -227,8 +227,19 @@ class LocalTensorSharedNumpyDataFunctor { DataType data_type = JUST(numpy::GetOFDataTypeFromNpArray(array)); Symbol device = JUST(Device::New("cpu")); const npy_intp* stride_ptr = PyArray_STRIDES(array); - const auto stride = std::make_shared(DimVector(stride_ptr, stride_ptr + dim)); - auto tensor_meta = std::make_shared(shape, data_type, device, stride, 0); + // stride + auto strides_vec = DimVector(stride_ptr, stride_ptr + dim); + auto element_size_in_bytes = PyArray_ITEMSIZE(array); + // NumPy strides use bytes. OneFlow strides use element counts. + for (auto& stride : strides_vec) { + if (stride % element_size_in_bytes != 0) { + return Error::RuntimeError() << "given numpy array strides not a multiple of the element " + "byte size. Copy the numpy array to reallocate the memory."; + } + stride /= element_size_in_bytes; + } + const auto strides = std::make_shared(strides_vec); + auto tensor_meta = std::make_shared(shape, data_type, device, strides, 0); // Build TensorBuffer const auto& Free = [obj](char* dptr) { diff --git a/python/oneflow/test/modules/test_from_numpy.py b/python/oneflow/test/modules/test_from_numpy.py index c4ef797b241..8d469c02ecf 100644 --- a/python/oneflow/test/modules/test_from_numpy.py +++ b/python/oneflow/test/modules/test_from_numpy.py @@ -28,6 +28,9 @@ def test_same_data(test_case): np_arr = np.random.randn(3, 4, 5) tensor = flow.from_numpy(np_arr) test_case.assertTrue(np.array_equal(np_arr, tensor.numpy())) + test_case.assertEqual(tensor.size(), (3, 4, 5)) + test_case.assertEqual(tensor.stride(), (20, 5, 1)) + test_case.assertEqual(tensor.storage_offset(), 0) np_arr[1:2, 2:3, 3:4] = random.random() test_case.assertTrue(np.array_equal(np_arr, tensor.numpy()))