U_NET 网络onnx部署

第一步是导出ONNX 

import os
import cv2
import archs_true
import onnxruntime
import torch
from albumentations import Compose
from albumentations.augmentations import transforms
from torch.utils.data import DataLoader





def pth_2onnx():
    # 创建 UNet 模型
    model = archs_true.UNet(num_classes=4)

    # 导入模型参数
    model_state_dict = torch.load("models/turbine_unet_random_15/model1000epoch_e-5.pth",
                                  map_location=lambda storage, loc: storage)
    model.load_state_dict(model_state_dict)

    # 将模型转换为 CPU 格式
    device = torch.device("cpu")
    model.to(device)
    model.eval()  # 进入评估模式

    # 创建虚拟输入（dummy input）
    dummy_input = torch.randn(1, 3, 800, 800, device=device)

    # 定义输入和输出的名称
    input_names = ["input"]
    output_names = ["output"]

    # 将模型转换为 ONNX 格式
    torch.onnx.export(model, dummy_input, "model.onnx", opset_version=12,verbose=False,  do_constant_folding=False, input_names=input_names,
                      output_names=output_names)

if __name__ == '__main__':
    pth_2onnx()

第二步进行C++部署

#include <fstream>
#include <iostream>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <onnxruntime_cxx_api.h>
#include <opencv2/core/utils/logger.hpp>
#include <opencv2/opencv.hpp>
#include<opencv2/dnn/dnn.hpp>
#include <algorithm>
#include "torch/script.h"
#include "torch/torch.h"

using namespace std;
using namespace cv;

//输入数据的长宽、通道数
static const int INPUT_H = 800;
static const int INPUT_W = 800;
static const int INPUT_CHANNELS =  3;
static const int OUTPUT_SIZE = 3 * 800 * 800;

int findMax(int a, int b, int c, int d) {
    int index = 0;
    int max_val = a;
    if (b > max_val) {
        max_val = b;
        index = 1;
    }
    if (c > max_val) {
        max_val = c;
        index = 2;
    }
    if (d > max_val) {
        max_val = d;
        index = 3;
    }
    return index;
}


cv::Mat flattenImage(const cv::Mat& image) {

    // 将图像 reshape 成一维向量
    cv::Mat flattened = image.reshape(1, image.rows * image.cols * 3 );

    // 转置矩阵，使得每个像素值都在一行中
    flattened = flattened.t();

    return flattened;
}

void test();

int main() {
    std::string imgpath = "121.png";
    cv::Mat test_image = cv::imread(imgpath);
    //cout << test_image;
    //cout << test_image << endl;
    cv::Scalar mean, std;
    mean[0] = 0.485 * 255;
    mean[1] = 0.456 * 255;
    mean[2] = 0.406 * 255;
    std[0] = 0.229 * 255;
    std[1] = 0.224 * 255;
    std[2] = 0.225 * 255;
    //cv::meanStdDev(test_image, mean, std);
    //cout << flattenedImage.size() << endl;
    //cout << test_image.channels() << endl;
    cout << mean << std << endl;
    //float m_mean = mean[0];
    //float m_std = std[0];
    utils::logging::setLogLevel(utils::logging::LOG_LEVEL_ERROR);//设置OpenCV只输出错误日志
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "U_NET");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);//设置线程数
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);//启用模型优化策略

    //CUDA option set
    OrtCUDAProviderOptions cuda_option;
    cuda_option.device_id = 0;
    cuda_option.arena_extend_strategy = 0;
    cuda_option.cudnn_conv_algo_search = OrtCudnnConvAlgoSearchExhaustive;
    cuda_option.gpu_mem_limit = SIZE_MAX;
    cuda_option.do_copy_in_default_stream = 1;
    //CUDA 加速
    //session_options.SetGraphOptimizationLevel() 函数用于设置在使用 ORT 库执行模型时应用的图优化级别。
    //ORT_ENABLE_ALL 选项启用所有可用的优化，这可以导致更快速和更高效的执行。
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
    //OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
    session_options.AppendExecutionProvider_CUDA(cuda_option);

#ifdef _WIN32//预处理器指令，用于在编译源代码时判断当前编译环境是否是WIndows环境
    /*
    wchar_t和普通的char之间的区别比较大。主要有以下几个方面：
    1.存储大小：wchar_t类型通常占用两个字节，而char类型通常占用一个字节。这是因为wchar_t类型用于表示Unicode字符集中的字符，
    每个字符需要占用更多的存储空间。
    2.字符集：wchar_t类型用于表示Unicode字符集中的字符，而char类型则用于表示ASCII字符集中的字符。
    Unicode字符集可以表示更广泛的字符集，包括中文、日文、韩文等，而ASCII字符集只能表示英文字母、数字和少量的符号。
    3.字符串表示方式：由于wchar_t类型占用更多的存储空间，
    因此在表示字符串时通常需要使用特殊的字符串类型，如const wchar_t*表示一个指向宽字符字符串的常量指针，
    而普通的char*则表示一个指向普通字符串的指针。
    4.字符串处理函数：由于wchar_t类型和普通的char类型在存储大小和字符集上有所不同，
    因此在处理字符串时通常需要使用不同的字符串处理函数。例如，wcslen()函数用于计算宽字符字符串的长度，
    而strlen()函数用于计算普通字符串的长度。
    */
    const wchar_t* model_path = L"model.onnx"; //表示一个指向常量宽字符字符串L"U_Net.onnx"的指针
#else
    const char* model_path = "yolov5s.onnx";
#endif

    // 第一个参数env是一个ORT环境对象，用于管理ORT会话的资源和配置。
    // 第二个参数model_path是一个指向ONNX模型文件的路径的常量指针，用于指定ORT会话要加载的模型。
    // 第三个参数session_options是一个ORT会话选项对象，用于配置ORT会话的选项和优化策略。
    Ort::Session session(env, model_path, session_options);
    // print model input layer (node names, types, shape etc.)
    Ort::AllocatorWithDefaultOptions allocator;//通过使用Ort::AllocatorWithDefaultOptions类，可以方便地进行内存分配和管理，而无需手动实现内存分配和释放的细节。

    // print number of model input nodes
    size_t num_input_nodes = session.GetInputCount();
    // 输入输出节点的名称
    std::vector<const char*> input_node_names = { "input" };
    std::vector<const char*> output_node_names = { "output" };

    //输出数据的大小
    const size_t input_tensor_size = 1 * INPUT_CHANNELS * INPUT_H * INPUT_W;

    //输入张量
    /*
        目前对于数据预处理方面有一些问题，首先是进行正则化，正则化之后除以255进行归一化
    */

    std::vector<float> input_tensor_values(input_tensor_size);
    //先把第一行的B提取完再将第二行的G
    for (int c = 0; c < 3; c++)
    {
        for (int i = 0; i < 800; i++)
        {
            for (int j = 0; j < 800; j++)
            {
                float pix = test_image.ptr<uchar>(i)[j * 3 +  c];//转换通道，输入onnx模型的图片通道顺序是RGB，但是opencv存储默认是BGR
                pix = (pix - mean[c]) / std[c];
                input_tensor_values[c * 800 * 800 + i * 800 + size_t(j)] = pix / 255.0;//归一化
                //cout << pix / 255.0 << endl;
            }
        }
    }
    //cv::Mat srcimg;
    //cv::cvtColor(test_image, test_image, cv::COLOR_BGR2RGB);
    //test_image.convertTo(srcimg, CV_32F, 1.0 / 255.0); // 将像素值转换为范围在0到1之间的浮点数
    //Mat blob = dnn::blobFromImage(test_image, 1.0 / 255, Size(800,800), Scalar(), true);//true代表进行翻转
    //输入张量的几个维度
    std::vector<int64_t> input_node_dims = { 1, 3, 800, 800 };

    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    //auto memory_info = Ort::MemoryInfo::CreateGpu(OrtDeviceAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(memory_info, input_tensor_values.data(), input_tensor_size, input_node_dims.data(), input_node_dims.size());

    std::vector<Ort::Value> ort_inputs;
    ort_inputs.push_back(std::move(input_tensor));//右值引用，避免不必要的拷贝和内存分配操作

    std::vector<Ort::Value> output_tensors;

    output_tensors = session.Run(Ort::RunOptions{ nullptr }, input_node_names.data(), ort_inputs.data(), input_node_names.size(), output_node_names.data(), output_node_names.size());

    // Get pointer to output tensor float values
    const float* rawOutput = output_tensors[0].GetTensorData<float>();
    //generate proposals
    std::vector<int64_t> outputShape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();

    //通过输出数据构造张量
    torch::Tensor tensor = torch::from_blob((void*)rawOutput, { 1, 4, 800, 800 });
    int dim = 1;
    torch::Tensor output2 = torch::softmax(tensor, dim);
    //已经正确
    //torch::print(output2);
    torch::Tensor _, output_max;
    std::tie(_, output_max) = torch::max(output2, dim);
    //torch::print(output_max);
    std::vector<int64_t> shape1 = output_max.sizes().vec();
    std::cout << shape1[0] << " " << shape1[1] << " " << shape1[2];
    //torch::Tensor output_visual_unique = output_max.unique();
    torch::Tensor out1, out2,out3;
    
    //判断各类0 2 3 的数目经比对几乎无差别
    std::tie(out1, out2, out3) = torch::_unique2(output_max,true,false, true);
    //torch::print(out1);
    //torch::print(out2);
    torch::print(out3);

    torch::Tensor output_visual1 = torch::zeros({ 1, 800, 800 }, torch::kCPU);
    output_visual1.masked_fill_(output_max == torch::Scalar(3), torch::Scalar(255.0));

    torch::Tensor output_visual2 = torch::zeros({ 1, 800, 800 }, torch::kCPU);
    output_visual2.masked_fill_(output_max == torch::Scalar(2), torch::Scalar(255.0));

    torch::Tensor output_visual3 = torch::zeros({ 1, 800, 800 }, torch::kCPU);
    output_visual3.masked_fill_(output_max == torch::Scalar(1), torch::Scalar(255.0));

    std::vector<torch::Tensor> output_visual_tensors = { output_visual1, output_visual2, output_visual3 };
    torch::Tensor output_visual = torch::cat(output_visual_tensors, 0);

    //进行转置
    torch::Tensor output_visual_transposed = output_visual.permute({ 1, 2, 0 });
    //变成连续存储
    torch::Tensor contiguous_tensor = output_visual_transposed.contiguous();
    // 获取张量的数据指针和形状信息
    // 注意Mat和tensor在内存中不一定是连续存储的，Mat一般是行连续存储
    float* data_ptr = contiguous_tensor.data_ptr<float>();
    std::vector<int64_t> shape = contiguous_tensor.sizes().vec();
    //Mat dstimg1(800, 800, CV_8UC3);
    //int num = 0;
    //for (int64_t i = 0; i < shape[0]; ++i) {
    //    for (int64_t j = 0; j < shape[1]; ++j) {
    //        for (int64_t k = 0; k < shape[2]; ++k) {
    //            // 使用索引访问元素
    //            float value = data_ptr[i * shape[1]*shape[2] + j * shape[2] + k];
    //            if (value == 255)
    //            {
    //                dstimg1.at<cv::Vec3b>(i, j)[k] = 255;
    //            }
    //        }
    //    }
    //}
    //cout << num << endl;
    //cv::imwrite("dst.jpg", dstimg1);
    //数据维度是对的800 * 800 * 3
    std::cout << static_cast<int>(shape[0]) << "    " << static_cast<int>(shape[1])<< "  " << static_cast<int>(shape[2]) << std::endl;
    // 创建 OpenCV Mat 对象，并从数据指针恢复图像
    cv::Mat image(static_cast<int>(shape[0]), static_cast<int>(shape[1]), CV_32FC3, data_ptr);
    // 将图像转换为 8 位无符号整数类型
    cv::Mat image_uint8;
    image.convertTo(image_uint8, CV_8UC3);
    cv::imwrite("resnew.png", image_uint8);
    system("pause");
    return 0;
}

void test() {
    torch::Tensor input = torch::rand({ 2, 3 });
    int dim = 1;
    torch::Tensor output = torch::softmax(input, dim);
}

这里边主要踩坑点，第一个是Mat和libtorch中的Tensor并非连续存储，进行数据转换的时候一定要注意，第二个是转换RRRGGGBBB的时候进行了正则化，对每个通道进行单独计算，可以采用opencv的库函数进行，libtorch的部署参考网上教程，注意添加环境变量后记得重启生效，不重启一般报错缺少dll