Refactor to numpy

oneflow/api/python/framework/tensor.cpp

@@ -49,6 +49,18 @@ const Symbol<DType>* GetTensorDType(const Tensor& tensor) {
   return &CHECK_JUST(DType::Get(tensor.dtype()->data_type()));
 }
 
+py::array ApiSwitchEagerTensorToNumpy(const py::handle& py_tensor) {
+  const std::shared_ptr<Tensor> tensor = py::cast<const std::shared_ptr<Tensor>>(py_tensor);
+  DataType data_type = tensor->dtype()->data_type();
+  switch (data_type) {
+#define EAGER_TENSOR_TO_NUMPY(cpp_type, of_type) \
+  case of_type: return EagerTensorToNumpy<cpp_type>(py_tensor).GetOrThrow();
+    OF_PP_FOR_EACH_TUPLE(EAGER_TENSOR_TO_NUMPY, POD_DATA_TYPE_SEQ BOOL_DATA_TYPE_SEQ)
+    default:
+      return Maybe<py::array>(Error::UnimplementedError() << "not support datatype").GetOrThrow();
+  }
+}
+
 void ApiEagerMirroredTensorZeros(const std::shared_ptr<Tensor>& tensor) {
   return EagerMirroredTensorZeros(tensor).GetOrThrow();
 }
@@ -198,6 +210,7 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
            [](const std::shared_ptr<one::Tensor>& tensor) {
              return CheckMetaConsistency(tensor).GetOrThrow();
            })
+      .def("to_numpy", &ApiSwitchEagerTensorToNumpy, py::return_value_policy::move)
 #define DEFINE_TENSOR_METHOD(T, type_proto)                    \
   .def("_copy_to_numpy_" #T, &ApiCopyMirroredTensorToNumpy<T>) \
       .def("_copy_from_numpy_" #T, &ApiCopyMirroredTensorFromNumpy<T>)

oneflow/api/python/utils/tensor_utils.h

@@ -30,14 +30,71 @@ limitations under the License.
 #include "oneflow/core/framework/tensor.h"
 #include "oneflow/core/framework/nd_sbp.h"
 #include "oneflow/core/functional/functional_api.yaml.h"
-
+#include "oneflow/core/framework/stride.h"
+#include "oneflow/core/register/ofblob.h"
+#include "oneflow/extension/python/numpy.h"
 namespace py = pybind11;
 
 namespace oneflow {
 namespace one {
 
 Maybe<void> EagerMirroredTensorZeros(const std::shared_ptr<Tensor>& t);
 
+template<typename T>
+inline static Maybe<py::array> EagerTensorToNumpy(const py::handle& py_tensor) {
+  const std::shared_ptr<Tensor> t = py::cast<const std::shared_ptr<Tensor>>(py_tensor);
+  py::handle handle;
+  std::shared_ptr<MirroredTensor> tensor;
+  CHECK_OR_RETURN(JUST(t->device()) == JUST(Device::New("cpu")));
+  CHECK_OR_RETURN(t->is_eager()) << "eager tensors supported only";
+  if (t->is_local()) {
+    tensor = JUST(t->AsMirroredTensor());
+    // set base object attr
+    handle = py::handle(py_tensor.ptr());
+  } else {
+    const Symbol<ConsistentTensorMeta>& tensor_meta = JUST(t->consistent_tensor_meta());
+    const Symbol<cfg::NdSbp>& nd_sbp = tensor_meta->nd_sbp();
+    CHECK_OR_RETURN(!nd_sbp->sbp_parallel().empty());
+    cfg::SbpParallel broadcast_sbp;
+    broadcast_sbp.mutable_broadcast_parallel();
+    std::vector<Symbol<cfg::SbpParallel>> sbp_tuple(nd_sbp->sbp_parallel_size(),
+                                                    SymbolOf(broadcast_sbp));
+    std::vector<Symbol<cfg::SbpParallel>> none;
+    const auto& consistent_tensor =
+        JUST(functional::ToConsistent(t, tensor_meta->parallel_desc(), sbp_tuple, none));
+    tensor = JUST(consistent_tensor->cur_rank_phy_tensor());
+  }
+
+  const size_t ndim = tensor->ndim();
+  const auto shape = numpy::OFShapeToNumpyShape(tensor->shape()->dim_vec());
+  // NumPy strides use bytes. OneFlow strides use element counts.
+  const auto stride = numpy::OFStrideToNumpyStride(JUST(tensor->stride())->StrideVec(),
+                                                   tensor->dtype()->data_type());
+
+  T* data_ptr = nullptr;
+  const auto& Callback = std::make_shared<std::function<void(uint64_t)>>([&](uint64_t ofblob_ptr) {
+    data_ptr = reinterpret_cast<OfBlob*>(ofblob_ptr)->mut_blob()->mut_dptr<T>();
+  });
+  bool is_printed = false;
+  SpinCounter::SpinWait(
+      1,
+      [&](const std::shared_ptr<SpinCounter>& sc) -> Maybe<void> {
+        return PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
+          return builder->SyncAccessBlobByCallback(tensor, sc, Callback, "mut");
+        });
+      },
+      [&is_printed]() {
+        if (!is_printed) {
+          blocking::StackInfoCallback();
+          is_printed = true;
+        }
+      });
+
+  return py::array(
+      py::buffer_info(data_ptr, sizeof(T), py::format_descriptor<T>::format(), ndim, shape, stride),
+      handle);
+}
+
 template<typename T>
 inline Maybe<void> CopyBetweenMirroredTensorAndNumpy(
     const std::shared_ptr<Tensor>& t, PyObject* array,

oneflow/extension/python/numpy.cpp

@@ -71,6 +71,22 @@ Maybe<DataType> GetOFDataTypeFromNpArray(PyArrayObject* array) {
   return NumpyTypeToOFDataType(np_array_type);
 }
 
+std::vector<size_t> OFShapeToNumpyShape(const DimVector& fixed_vec) {
+  size_t ndim = fixed_vec.size();
+  auto result = std::vector<size_t>(ndim);
+  for (int i = 0; i < ndim; i++) { result[i] = fixed_vec.at(i); }
+  return result;
+}
+
+// NumPy strides use bytes. OneFlow strides use element counts.
+std::vector<size_t> OFStrideToNumpyStride(const StrideVector& fixed_vec, const DataType data_type) {
+  size_t ndim = fixed_vec.size();
+  auto result = std::vector<size_t>(ndim);
+  int byte_per_elem = GetSizeOfDataType(data_type);
+  for (int i = 0; i < ndim; i++) { result[i] = fixed_vec.at(i) * byte_per_elem; }
+  return result;
+}
+
 // Executing any numpy c api before _import_array() results in segfault
 // NOTE: this InitNumpyCAPI() works because of `PY_ARRAY_UNIQUE_SYMBOL`
 // defined in numpy_internal.h

oneflow/extension/python/numpy_internal.h

@@ -22,6 +22,8 @@ limitations under the License.
 // ************************
 
 #include "oneflow/core/common/data_type.h"
+#include "oneflow/core/common/fixed_vector.h"
+#include "oneflow/core/common/shape_vec.h"
 
 // PyArrayObject cannot be forward declared, or a compile error will occur
 
@@ -56,6 +58,10 @@ Maybe<DataType> NumpyTypeToOFDataType(int np_array_type);
 
 Maybe<DataType> GetOFDataTypeFromNpArray(PyArrayObject* array);
 
+std::vector<size_t> OFShapeToNumpyShape(const DimVector& fixed_vec);
+
+std::vector<size_t> OFStrideToNumpyStride(const StrideVector& fixed_vec, const DataType data_type);
+
 Maybe<void> InitNumpyCAPI();
 
 }  // namespace numpy

python/oneflow/framework/docstr/tensor.py

@@ -523,6 +523,15 @@
     """,
 )
 
+add_docstr(
+    oneflow.Tensor.cpu,
+    """
+    Tensor.cpu() → Tensor
+
+    Returns a copy of this object in CPU memory. if this object is already in CPU memory and on the correct device, then no copy is performed and the original object is returned.
+    """,
+)
+
 add_docstr(
     oneflow.Tensor.pow,
     """

python/oneflow/framework/tensor.py

@@ -29,29 +29,6 @@
 TensorTuple = flow._oneflow_internal.TensorTuple
 
 
-def _tensor_numpy(eager_local_tensor):
-    assert (
-        not eager_local_tensor.is_lazy
-    ), "tensor.numpy() is not allowed to called in nn.Graph.build(*args) or called by lazy tensor."
-    if eager_local_tensor.dtype == flow.tensor_buffer:
-        shapes, dtypes = eager_local_tensor._tensor_buffer_shapes_and_dtypes
-        tensors = flow.tensor_buffer_to_list_of_tensors(
-            eager_local_tensor, shapes, dtypes
-        )
-        return [t.numpy() for t in tensors]
-    method_name = eager_local_tensor._get_copy_mirrored_tensor_to_numpy_func_name()
-    copy_to_numpy = getattr(eager_local_tensor, method_name)
-
-    ndarray = np.empty(
-        shape=tuple(eager_local_tensor.shape),
-        dtype=flow.convert_oneflow_dtype_to_numpy_dtype(eager_local_tensor.dtype),
-    )
-
-    if ndarray.size != 0:
-        copy_to_numpy(ndarray)
-    return ndarray
-
-
 def _size(self, idx=None):
     if idx is None:
         return self.shape
@@ -733,7 +710,7 @@ def _get_device(self):
 
 def _format(self, format_spec):
     if self.dim() == 0:
-        return self.numpy().tolist().__format__(format_spec)
+        return self.tolist().__format__(format_spec)
     return object.__format__(self, format_spec)
 
 
@@ -745,13 +722,27 @@ def _gather(self, dim, index):
     return flow._C.dim_gather(self, dim, index, False)
 
 
+def _cpu(self):
+    if self.device == flow.device("cpu"):
+        return self
+    return self.to(device="cpu")
+
+
+def _numpy(self):
+    assert self.device == flow.device("cpu")
+    assert (
+        not self.is_lazy
+    ), "tensor.numpy() is not allowed to called in nn.Graph.build(*args) or called by lazy tensor."
+    return self.to_numpy()
+
+
 def RegisterMethods():
     Tensor.__mul__ = lambda self, other: self.mul(other)
     Tensor.__rmul__ = lambda self, other: self.mul(other)
     Tensor.__add__ = lambda self, other: self.add(other)
     Tensor.__iadd__ = lambda self, other: self.add_(other)
     Tensor.ndim = property(_ndim)
-    Tensor.numpy = _tensor_numpy
+    Tensor.numpy = _numpy
     Tensor.size = _size
     Tensor.dim = _ndim
     Tensor.ndimension = _ndim
@@ -893,6 +884,7 @@ def RegisterMethods():
     Tensor.gather = _gather
     Tensor.all = _all
     Tensor.any = _any
+    Tensor.cpu = _cpu
 
 
 def register_tensor_op(op_name):

python/oneflow/nn/modules/tensor_ops.py

@@ -221,6 +221,8 @@ def item_op(input):
         1.0
     """
     assert input.numel() == 1, "Only a Tensor with 1 element can be converted to Scalar"
+    if input.device != flow.device("cpu"):
+        input = input.cpu()
     return input.numpy().item()
 
 
@@ -245,6 +247,8 @@ def tolist_op(input):
     """
     if input.numel() == 1 and input.ndim == 0:
         return input.item()
+    if input.device != flow.device("cpu"):
+        input = input.to(device="cpu")
     return input.numpy().tolist()