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

posted @ 2022-01-16 17:19  Keep_Silent  阅读(128)  评论(0编辑  收藏  举报