pytorch学习004- -libtorch处理多输入/多输出问题

一般处理多输入/多输出问题时，pytorch中容易处理；但是在libtorch中会出现一些问题。
这篇博客为记录用。

准备工作

首先加载多模型。
因为可能需要根据不同情况调用多个模型，所以这里预先声明一个models

std::vector<torch::jit::script::Module> models;

然后使用try/catch来加载模型

try
      {
        models.push_back(torch::jit::load("xxxxxx.pt"));

        models.push_back(torch::jit::load("xxxxxx.pt"));

        models.push_back(torch::jit::load("xxxxxx.pt"));

        models.push_back(torch::jit::load("xxxxxx.pt"));
      }
      catch (const c10::Error& e)
      {
        std::cerr << "error loading models\n";
      }

多输入问题

一种方案是调整网络的输入，在pytorch中进行修改；另一种方案是在libtorch中修改。

调整网络模型

xxx

设置libtorch输入

定义一个vector

std::vector<torch::IValue> all_input;
  all_input.push_back(a);
  all_input.push_back(b);
  all_input.push_back(c);

多输出问题

...
auto outputs = module->forward(inputs).toTuple();
torch::Tensor out1 = outputs->elements()[0].toTensor();
torch::Tensor out2 = outputs->elements()[1].toTensor();
...

参考文章

知乎
大专栏
github