将GpuMat复制到CUDA张量

Question

我试图在C++中运行模型推理。

我成功地用torch.jit.trace跟踪了Python中的模型。

我能够在C++中使用torch::jit::load()加载模型。

我能够在cpu和gpu上进行推理，但是起点总是torch::from_blob方法，它似乎正在创建cpu端张量。

为了提高效率，我想直接向CUDA张量转换/复制cv::cuda::GpuMat。为了寻找简单的例子，我一直在测试和文档中挖掘，但却找不到。

问:如何从cv:: CUDA ::GpuMat创建cuda张量？

Answer 1

下面是一个示例：

//define the deleter ...
void deleter(void* arg) {};

//your convert function 

cuda::GpuMat gImage;

//build or load your image here ...

std::vector<int64_t> sizes = {1, static_cast<int64_t>(gImage.channels()),
                          static_cast<int64_t>(gImage.rows),
                          static_cast<int64_t>(gImage.cols)};

long long step = gImage.step / sizeof(float);

std::vector<int64_t> strides = {1, 1, step, static_cast<int64_t>(gImage.channels())};

auto tensor_image = torch::from_blob(gImage.data, sizes, strides, deleter,  torch::kCUDA);
std::cout << "output tensor image : " << tensor_image << std::endl;