如何将 CuPy 数组传输到 tensorflow

Posted 2023-02-16

【中文标题】如何将 CuPy 数组传输到 tensorflow

【英文标题】：How to transfer CuPy arrays to tensorflow

【发布时间】：2019-11-13 15:49:04

【问题描述】：

我正在使用 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 ？

【参考方案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))

【参考方案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

【讨论】：

似乎 CuPy 有一个特殊的 PyTorch API，允许在 GPU 上将 CuPy 数组转换为 PyTorch 张量，而无需通过 CPU 上的 NumPy。但是，缺少对 TensorFlow 的这种支持 :-(

CuPy 支持标准协议（DLPack 和 cuda_array_interface），但 TF 不支持。似乎正在进行一些讨论：github.com/tensorflow/tensorflow/issues/24453github.com/tensorflow/tensorflow/issues/29039