Я пытаюсь передать вывод некоторой операции 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)
Однако я хотел бы избежать накладных расходов при копировании памяти с устройства на хост и с хоста на устройство.
Я не смог найти способ построить mxnet.ndarray.NDArray непосредственно из h_output.
Самая близкая вещь, которую мне удалось найти - это конструкция ndarray из dlpack .
Но не понятно, как работать с объектом dlpack из python.
Есть ли способ достичь взаимодействия NDArray <-> pycuda без копирования памяти через хост?