TensorRT部署U_Net

一、U_NET  ONNX部署

1. U_NET 模型导出为onnx

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)

2. 数据的预处理问题

数据预处理主要包括两个过程，首先是正则化，之后所有数据除以255，再进行维度转换，将（800,800,3）转化为（3,800,800）

//4维向量     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;     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;//归一化             }         }     }

3. 数据的后处理问题

数据后处理主要是将输出的结果（1,4,800,800）的张量通过对每个点的四个特征值的大小进行比较，找出最大值，0代表无缺陷，1带边第一类缺陷，2代表第二类缺陷，3代表第三类缺陷，同时建立他们之间的映射关系，0映射黑色，1映射红色….。对于一些张量操作，使用torch版本的C++API库Libtorch进行，目前还不清楚能否在没有pytorch的设备上运行，下周主要尝试将代码转移到检测电脑上。

//通过输出数据构造张量 torch::Tensor tensor = torch::from_blob((void*)rawOutput, { 1, 4, 800, 800 }); int dim = 1; torch::Tensor output2 = torch::softmax(tensor, dim); torch::Tensor _, output_max; std::tie(_, output_max) = torch::max(output2, dim); std::vector<int64_t> shape1 = output_max.sizes().vec(); 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 }); //将Tensor变成连续存储 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(); // 创建 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);

二、U_NET  TensorRT部署

1. U_NET模型导出为engine

#include <fstream> #include <iostream> #include <sstream> #include "NvInfer.h" #include "NvOnnxParser.h" #include "NvinferRuntime.h" using namespace nvinfer1; using namespace nvonnxparser; // 全局创建 ILogger 类型的对象 class Logger : public ILogger {     virtual void log(Severity severity, const char* msg) noexcept override     {         // suppress info-level messages         if (severity != Severity::kINFO)             std::cout << msg << std::endl;     } } gLogger; int onnx2engine(std::string onnx_filename, std::string enginefilePath, int type) {     // 创建builder     IBuilder* builder = createInferBuilder(gLogger);     // 创建network     nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));     // 创建onnx模型解析器     auto parser = nvonnxparser::createParser(*network, gLogger);     // 解析模型     parser->parseFromFile(onnx_filename.c_str(), 2);     for (int i = 0; i < parser->getNbErrors(); ++i)     {         std::cout << parser->getError(i)->desc() << std::endl;     }     printf("tensorRT load onnx model sucessful! \n"); // 解析模型成功，第一个断点测试     // 使用builder对象构建engine     IBuilderConfig* config = builder->createBuilderConfig();     config->setMaxWorkspaceSize(16 * (1 << 20));  // 设置最大工作空间     config->setFlag(BuilderFlag::kGPU_FALLBACK);  // 启用GPU回退模式，作用？     //  config->setFlag(BuilderFlag::kSTRICT_TYPES);  //强制执行xx位的精度计算     if (type == 1) {         config->setFlag(nvinfer1::BuilderFlag::kFP16); // 设置精度计算     }     if (type == 2) {         config->setFlag(nvinfer1::BuilderFlag::kINT8);     }     //IOptimizationProfile* profile = builder->createOptimizationProfile(); //创建优化配置文件     //profile->setDimensions("x", OptProfileSelector::kMIN, Dims4(1, 3, 32, 300)); // 设置输入x的动态维度，最小值     //profile->setDimensions("x", OptProfileSelector::kOPT, Dims4(1, 3, 32, 320)); // 期望输入的最优值     //profile->setDimensions("x", OptProfileSelector::kMAX, Dims4(1, 3, 32, 340)); // 最大值     //config->addOptimizationProfile(profile);     ICudaEngine* myengine = builder->buildEngineWithConfig(*network, *config); //创建engine  第二个断点测试     std::cout << "try to save engine file now" << std::endl;     std::ofstream p(enginefilePath, std::ios::binary);     if (!p) {         std::cerr << "could not open plan output file" << std::endl;         return 0;     }     // 序列化     IHostMemory* modelStream = myengine->serialize(); // 第三个断点测试     p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size()); // 写入     modelStream->destroy(); // 销毁     myengine->destroy();     network->destroy();     parser->destroy();     std::cout << "convert onnx model to TensorRT engine model successfully!" << std::endl; // 转换成功，第四个断点测试     return 0; } int main(int argc, char** argv) {     onnx2engine("model.onnx", "model.engine", 1);     return 0; }

可以导出FP16也可以选择导出FP32，计算结果差别不大，模型文件大小差别FP32是FP16的2倍。

2.数据预处理

采用CUDA实现并行处理，相比于之前通过for循环串行处理时间缩短5ms。

• 串行处理代码:

std::vector<float> input_tensor_values(input_data_length); clock_t start_time = clock(); //预处理代码使用cuda实现试试 //先把第一行的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;//归一化         }     } } cudaError_t err3 = cudaMemcpyAsync(data_buffer[input_node_index], input_tensor_values.data(), input_data_length * sizeof(float), cudaMemcpyHostToDevice, stream); if (err3 != cudaSuccess) {     std::cout << "Failed to transfer input data to GPU: " << cudaGetErrorString(err3) << std::endl;     return -1; }

• 并行处理

//均值 struct mean {     float value[3]; }; //标准差 struct standard_deviation {     float value[3]; }; //进行预处理的核函数 __global__ void process_kernel(uint8_t* src, int src_line_size, int src_width,     int src_height, float* dst, standard_deviation st, mean me, int edge){     int position = blockDim.x * blockIdx.x + threadIdx.x;     if (position >= edge) return;     int dx = position % src_width;     int dy = position / src_height;     float c0 = *(src + src_line_size * dy + dx * 3);     float c1 = *(src + src_line_size * dy + dx * 3 + 1);     float c2 = *(src + src_line_size * dy + dx * 3 + 2);         c0 = (c0 - me.value[0]) / st.value[0];     c1 = (c1 - me.value[1]) / st.value[1];     c2 = (c2 - me.value[2]) / st.value[2];     // normalization     c0 = c0 / 255.0f;     c1 = c1 / 255.0f;     c2 = c2 / 255.0f;     // rgbrgbrgb to rrrgggbbb，不需要×3，首先要知道在做什么     int area = src_width * src_height;     //C0 C1 C2 分别对应每个像素点的 B G R通道     float* pdst_c0 = dst + dy * src_width + dx;     float* pdst_c1 = pdst_c0 + area;     float* pdst_c2 = pdst_c1 + area;     *pdst_c0 = c0;     *pdst_c1 = c1;     *pdst_c2 = c2; } void cuda_preprocess(uint8_t* src, int src_width, int src_height,     float* dst, uint8_t* input_data_host, uint8_t* input_data_device,     cudaStream_t stream) {     int img_size = src_width * src_height * 3;     // copy data to pinned memory     memcpy(input_data_host, src, img_size);     // copy data to device memory     CUDA_CHECK(cudaMemcpyAsync(input_data_device, input_data_host, img_size, cudaMemcpyHostToDevice, stream));     //均值 标准差     mean me;     standard_deviation st;     me.value[0] = 0.485 * 255;     me.value[1] = 0.456 * 255;     me.value[2] = 0.406 * 255;     st.value[0] = 0.229 * 255;     st.value[1] = 0.224 * 255;     st.value[2] = 0.225 * 255;     int jobs = src_width * src_height;     int threads = 256;     int blocks = ceil(jobs / (float)threads);     //blocks是2500     //开始执行核函数     process_kernel << <blocks, threads, 0, stream >> > (input_data_device, src_width * 3, src_width,         src_height, dst, st, me, jobs); }

3.数据后处理

数据后处理这块主要遇到了两个问题，第一个是从显卡中转运数据时间太长，第二个是libtorch中矩阵张量的计算考虑转移到GPU进行，看能否缩短计算时间，经过测试并不行，涉及到libtorch中无法调用cuda的一些解决办法。

• 通过开辟页锁定内存加快数据传输的速度

1. 原方案：

//普通内存的指针，开辟在堆区，逻辑地址 float* result_array = new float[output_data_length]; //数据的后处理 cudaError_t err4 = cudaMemcpyAsync(result_array, data_buffer[output_node_index], output_data_length * sizeof(float), cudaMemcpyDeviceToHost, stream); if (err4 != cudaSuccess) {     std::cout << "Failed to transfer output data to HOST: " << cudaGetErrorString(err4) << std::endl;     return -1; }

• 新的解决方案

//页锁定内存---后处理 float* result_array = nullptr; cudaMallocHost((void**)&result_array, output_data_length * sizeof(float)); //数据的后处理 cudaError_t err4 = cudaMemcpyAsync(result_array, data_buffer[output_node_index], output_data_length * sizeof(float), cudaMemcpyDeviceToHost, stream); if (err4 != cudaSuccess) {     std::cout << "Failed to transfer output data to HOST: " << cudaGetErrorString(err4) << std::endl;     return -1; } //解析结果,libtorch无法调用GPU进行运算 //通过输出数据构造张量,时间较长，排查一下原因 torch::Tensor tensor = torch::from_blob((float*)result_array, { 1, 4, 800, 800 });

原来耗时大概30ms，现在耗时大概1ms

• Libtorch调用cuda进行加速需要注意的问题

1. 如果想要推理torch训练的模型注意下载对应版本libtorch。
2. Libtorch1.13.0可能遇到无法启动cuda情况torch::cuda::is_available()是0，需要在命令行（1.13.0）添加/INCLUDE:"?ignore_this_library_placeholder@@YAHXZ"

torch::Tensor output2;  // 输出张量 //将输入张量和输出张量移动到 GPU 上 torch::Tensor tensor1 = tensor.to(torch::kCUDA); torch::Tensor output3 = output2.to(torch::kCUDA); //在 GPU 上计算 softmax // 假设 dim 为 1 output2 = torch::softmax(tensor, 1); // 将输出张量移回到 CPU output2 = output2.to(torch::kCPU);

经测试时间差别几乎没有，可能是由于输出结果的张量数据量并不大。

4.TensorRT检测效果

处理一张图片的总时间：平均10ms