我已经训练了 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", …Run Code Online (Sandbox Code Playgroud)