如何将 CuPy 数组传输到 tensorflow

Question

我正在使用 CuPy 在 GPU 上生成一些地面实况 + 网络输入数据，我想通过标准 API 将其发送到 TensorFlow 网络：

feed_dict = {
    network.Labels: labels,
    network.Network_Input: networkInput
}

但是，TensorFlow 1.15 需要 NumPy 而不是 CuPy，我收到此错误：

File "/lib/python3.6/site-packages/numpy/core/numeric.py", line 538, in asarray return array(a, dtype, copy=False, order=order)
ValueError: object __array__ method not producing an array

尝试使用将 CuPy 数组转换为 NumPy 数组

labels = numpy.array(labels, dtype=np.float)

我得到错误：

ValueError: object __array__ method not producing an array

有没有办法将 CuPy 数据传递给 tensorflow ？

Answer 1

tf.experimental.dlpack 包中的 Tensorflow 2 现在支持 DLPack 张量结构。

你可以在这个repo中找到howto用法：https://github.com/VoVAllen/tf-dlpack#cupy---tensorflow

# pip install numba numpy
import numpy as np
from numba import cuda

# NumPy - CPU Array
cpu_arr = np.random.rand(10_000, 10_000)

# Use Numba to move to GPU
numba_gpu_arr = cuda.to_device(cpu_arr)

# Use CuPy's asarray function and toDlpack to create DLPack capsule. There are multiple other ways to do this (i.e. PyTorch Utils)
dlpack_arr = cp.asarray(numba_gpu_arr).toDlpack()

# Migrate from Numba, used for custom CUDA JIT kernels to PyTorch
tf_tensor = tfdlpack.from_dlpack(dlpack_arr)

# Confirm TF tensor is on GPU
print(tf_tensor.device)

# Use `tfdlpack` to migrate back to Numba
dlpack_capsule = tfdlpack.to_dlpack(tf_tensor)
numba_arr = cuda.to_device(cp.fromDlpack(dlpack_capsule))

Answer 2

要将 CuPy 数组（在 GPU 上）复制到 NumPy 数组（在 CPU 上），请使用 cupy.asnumpy(arr) 或 arr.get()。

https://docs-cupy.chainer.org/en/latest/reference/ndarray.html#conversion-to-from-numpy-arrays