如何在c++中提取割炬model的output？

Question

我已經接受了培訓 keras model 並使用mmdnn對其進行了轉換。 然后我嘗試在 c++ 代碼中使用它：

#include <iostream>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>

#include <torch.h>

int main()
{
    cv::Mat image;
    image= cv::imread("test_img.png", cv::IMREAD_GRAYSCALE);   // Read the file

try
{
    torch::jit::script::Module module;
    module = torch::jit::load("my_model.pth");

    torch::IntArrayRef input_dim = std::vector<int64_t>({ 1, 2, 256, 256});

    cv::Mat input_img;
    image.convertTo(input_img, CV_32FC3, 1 / 255.0);
    torch::Tensor x = torch::from_blob(input_img.data, { 1, 2, 256, 256 }, torch::kFloat);
    torch::NoGradGuard no_grad;

    auto output = module.forward({ x });

    float* data = static_cast<float*>(output.toTensor().data_ptr());

    cv::Mat output_img = cv::Mat(256, 256, CV_32FC3, data);
    cv::imwrite("output_img.png", output_img);
}
catch (std::exception &ex)
{
    std::cout << "exception! " << ex.what() << std::endl;
}

    return 0;
}

此代碼引發異常：

例外：isTensor() 內部斷言在 E.\20B\pytorch\pytorch\aten\src\ATen/core/ivalue_inl:h,112 處失敗。 請向 PyTorch 報告錯誤：Expected Tensor but got Tuple (toTensor at E.\20B\pytorch\pytorch\aten\src\ATen/core/ivalue_inl:h:112) (沒有可用的回溯)

這是在行中拋出的float* data = static_cast<float*>(output.toTensor().data_ptr()); 當調用 function toTensor()時。 如果我使用toTuple()而不是toTensor()則結果沒有 function data_ptr() ，但我需要它來提取數據（並將其放入 opencv 圖像）。

如何從 model output 中提取圖像？

Answer 1

在這種情況下，model 的答案是 2 個圖像的元組。 我們可以通過以下方式提取它們：

torch::Tensor t0 = output.toTuple()->elements()[0].toTensor();
torch::Tensor t1 = output.toTuple()->elements()[1].toTensor();

變量t0和t1包含具有 model 的 output 的張量。