如何在 C++ 中提取火炬模型的输出？

Posted 2023-03-29

【中文标题】如何在 C++ 中提取火炬模型的输出？

【英文标题】：How to extract output of torch model in c++?

【发布时间】：2020-09-21 07:43:31

【问题描述】：

我已经训练好 keras 模型并使用 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 报告错误。预期张量但得到元组（toTensor at E:\20B\pytorch\pytorch\aten\src\ATen/core/ivalue_inl.h:112）（无 回溯可用）

当函数toTensor() 被调用时，这是在float* data = static_cast<float*>(output.toTensor().data_ptr()); 行中抛出的。如果我使用toTuple() 而不是toTensor()，则结果没有data_ptr() 函数，但我需要它来提取数据（并将其放入opencv 图像）。

如何从模型输出中提取图像？

【参考方案1】：

在这种情况下，模型的答案是 2 个图像的元组。我们可以这样提取它们：

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

变量t0 和t1 包含带有模型输出的张量。