将内存 blob 分配给 py-torch 输出张量（C++ API）

Question

我正在使用Py-torch培训线性模型，并将其保存到带有“保存”功能调用的文件中。我还有另一个代码将模型加载到C ++中并执行推理。 我想指示火炬CPP库在最终输出张量上使用特定的内存斑点。这甚至可能吗？如果是，怎么样？在下面，您可以看到我要实现的一小部分。

#include <iostream>
#include <memory>

#include <torch/script.h>

int main(int argc, const char* argv[]) {
  if (argc != 3) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }
  long numElements = (1024*1024)/sizeof(float) * atoi(argv[2]);

  float *a = new float[numElements]; 
  float *b = new float[numElements];
  float *c = new float[numElements*4];

  for (int i = 0; i < numElements; i++){
    a[i] = i;
    b[i] = -i;
  }

  //auto options = torch::TensorOptions().dtype(torch::kFloat64);
  at::Tensor a_t = torch::from_blob((float*) a, {numElements,1});
  at::Tensor b_t = torch::from_blob((float*) b, {numElements,1});
  at::Tensor out = torch::from_blob((float*) c, {numElements,4});

  at::Tensor c_t = at::cat({a_t,b_t}, 1);
  at::Tensor d_t = at::reshape(c_t, {numElements,2});

  torch::jit::script::Module module;
  try {
    module = torch::jit::load(argv[1]);
  }
  catch (const c10::Error& e) {
    return -1;
  }


  out =  module.forward({d_t}).toTensor();
  std::cout<< out.sizes() << "\n";

  delete [] a;
  delete [] b;
  delete [] c;

  return 0;
}

因此，我将内存分配到“ C”中，然后从此内存中施放张量。我将此内存存储在名为“ Out”的张量中。当我调用正向方法时，我加载模型。我观察到所得数据被复制/移至“ OUT”张量。但是，我想指示火炬直接存储在“输出”内存中。这可能吗？

Answer 1

在 libtorch 源代码中的某个地方（我不记得在哪里，我会尝试找到该文件），有一个类似于下面的运算符（注意最后一个 &&），

torch::tensor& operator=(torch::Tensor rhs) &&;

如果我没记错的话，它可以满足您的需要。基本上，torch 假设如果您将张量 rhs 分配给右值引用张量，那么您实际上意味着将 rhs 复制到底层存储中。

所以在你的情况下，那就是

std::move(out) = module.forward({d_t}).toTensor();

或

torch::from_blob((float*) c, {numElements,4}) = module.forward({d_t}).toTensor();