python - 如何将 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 张量结构包。

您可以在此存储库中找到使用方法: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))