从 pycuda.driver.DeviceAllocation 创建 mxnet.ndarray.NDArray

Question

我正在尝试将一些 pycuda 操作的输出传递给 mxnet 计算图的输入。我可以使用以下代码通过 numpy 转换来实现此目的

import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import mxnet as mx

batch_shape = (1, 1, 10, 10)
h_input = np.zeros(shape=batch_shape, dtype=np.float32)
# init output with ones to see if contents really changed
h_output = np.ones(shape=batch_shape, dtype=np.float32)
device_ptr = cuda.mem_alloc(input.nbytes)
stream = cuda.Stream()
cuda.memcpy_htod_async(d_input, h_input, stream)

# here some actions with d_input may be performed, e.g. kernel calls
# but for the sake of simplicity we'll just transfer it back to host
cuda.memcpy_dtoh_async(d_input, h_output, stream)
stream.synchronize()
mx_input = mx.nd(h_output, ctx=mx.gpu(0))

print('output after pycuda calls: ', h_output)
print('mx_input: ', mx_input)

但是我想避免设备到主机和主机到设备内存复制的开销。

我找不到直接从 h_output 构建 mxnet.ndarray.NDArray 的方法。我能找到的最接近的东西是 ndarray from dlpack 的构造。但不清楚如何使用 python.

中的 dlpack 对象

有没有办法在不通过主机复制内存的情况下实现 NDArray <-> pycuda 互操作性？

Answer 1

很遗憾，目前无法做到。

从 pycuda.driver.DeviceAllocation 创建 mxnet.ndarray.NDArray

Create mxnet.ndarray.NDArray from pycuda.driver.DeviceAllocation

python

pycuda

mxnet