c++ - 将 GpuMat 复制到 CUDA 张量

Question

我正在尝试在 C++ 中运行模型推理。
我使用 torch.jit.trace 在 Python 中成功追踪了模型。
我可以使用 torch::jit::load() 在 C++ 中加载模型。
我能够在 cpu 和 gpu 上执行推理，但是起点始终是 torch::from_blob 方法，它似乎正在创建 cpu 端张量。
为了提高效率，我想将 cv::cuda::GpuMat 直接转换/复制到 CUDA Tensor。我一直在挖掘 pytorch tests和 docs正在寻找方便的示例，但无法找到。

问题:如何从 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;