fix_tensor_from_numpy_mem_leak_bug (#8391)

* fix_tensor_from_numpy_mem_leak_bug * add note * refine note * refine Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Oneflow-Inc · Jun 9, 2022 · d98132e · d98132e
1 parent 469f72d
commit d98132e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/oneflow/api/python/framework/tensor.cpp b/oneflow/api/python/framework/tensor.cpp
@@ -337,6 +337,7 @@ static PyObject* PyTensorObject_type(PyObject* self, PyObject* args, PyObject* k
     ASSERT(CopyBetweenMirroredTensorAndNumpy<T>(PyTensor_Unpack(self), copied,            \
                                                 BlobNumpyCopyUtil<T>::From, "mut",        \
                                                 /*block_host_until_done=*/false));        \
+    Py_DECREF(copied);                                                                    \
     Py_RETURN_NONE;                                                                       \
     END_HANDLE_ERRORS                                                                     \
   }

diff --git a/oneflow/api/python/functional/tensor_api.cpp b/oneflow/api/python/functional/tensor_api.cpp
@@ -239,7 +239,12 @@ class LocalTensorSharedNumpyDataFunctor {
     if (!PyArray_IS_C_CONTIGUOUS(array)) {
       OF_LOG_ONCE(LOG(WARNING) << "OneFlow don't support non-contiguous array now, "
                                   "and we will copy the array to a contiguous one.");
+      // PyArray_GETCONTIGUOUS will return a reference if array is already contiguous,
+      // otherwise return a (contiguous) copy of the array.
+      // Note: Increment the reference count for array occurs whether the array is continuous or not
       array = PyArray_GETCONTIGUOUS(array);
+    } else {
+      Py_INCREF(obj);
     }
 
     // Build TensorMeta
@@ -264,13 +269,12 @@ class LocalTensorSharedNumpyDataFunctor {
     auto tensor_meta = std::make_shared<MirroredTensorMeta>(shape, strides, data_type, device, 0);
 
     // Build TensorBuffer
-    const auto& Free = [obj](char* dptr) {
+    const auto& Free = [array](char* dptr) {
       CHECK_JUST(Global<ForeignLockHelper>::Get()->WithScopedAcquire([&]() -> Maybe<void> {
-        Py_DECREF(obj);
+        Py_DECREF(array);
         return Maybe<void>::Ok();
       }));
     };
-    Py_INCREF(obj);  // make TensorBuffer hold ndarray
     void* data_ptr = PyArray_DATA(array);
     auto array_size_in_bytes = PyArray_NBYTES(array);
     auto tensor_data = std::make_shared<vm::TensorStorage>();