Libtorch：如何使用 GPU 指针制作张量？

Question

下面是我想要做的伪代码。我已经知道如何将张量移动到 GPU ( .cuda())...
但不知道如何使用 GPU 指针来创建新的张量。有什么方法我错过了吗？
我不想复制devPtr回主机端，而只是用指针制作 GPU 张量。

int main(void) {
  float* devPtr;
  cudaMalloc((void**)&devPtr, sizeof(float)*HOSTDATA_SIZE);
  cudaMemcpy(devPtr, hostData, sizeof(float)*HOSTDATA_SIZE, cudaMemcpyHostToDevice);

  torch::Tensor inA = /* make Tensor with devPtr which is already in GPU */;
  torch::Tensor inB = torch::randn({1, 10, 512, 512}).cuda();

  torch::Tensor out = torch::matmul(inA, inB);

  std::cout << out << std::endl;
  return 0;
}

Answer 1

我认为这应该有效，你能确认吗？

auto dims = torch::IntArrayRef{1, 10, 512, 512};
auto gpu_tensor = torch::from_blob(dev_ptr, dims, torch::TensorOptions().device(torch::kCUDA))

请注意，torch::from_blob 不获取指针的所有权。如果您需要使 gpu_tensor 独立于 的生命周期dev_ptr，那么您需要克隆它。